Update app.py

app.py

@@ -1,36 +1,51 @@
 import gradio as gr
-import spaces
 from transformers import AutoModelForCausalLM, AutoProcessor
 import torch
 from PIL import Image
 
-# Load model and processor
+# Model and Processor Initialization
 models = {
-    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
+    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained(
+        "microsoft/Phi-3.5-vision-instruct", 
+        trust_remote_code=True, 
+        torch_dtype="auto", 
+        _attn_implementation="flash_attention_2"
+    ).cuda().eval()
 }
 
 processors = {
     "microsoft/Phi-3.5-vision-instruct": AutoProcessor.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True)
 }
 
-# Set the default query
-DEFAULT_QUERY = "You are an image to prompt converter. Your work is to observe each and every detail of the image and craft a detailed prompt under 100 words in this format: [image content/subject, description of action, state, and mood], [art form, style], [artist/photographer reference if needed], [additional settings such as camera and lens settings, lighting, colors, effects, texture, background, rendering]."
+# Default question
+default_question = (
+    "You are an image-to-prompt converter. Your work is to observe each and every detail of the image and "
+    "craft a detailed prompt under 100 words in this format: [image content/subject, description of action, state, "
+    "and mood], [art form, style], [artist/photographer reference if needed], [additional settings such as camera "
+    "and lens settings, lighting, colors, effects, texture, background, rendering]."
+)
 
-@spaces.GPU
-def run_example(image, model_id="microsoft/Phi-3.5-vision-instruct"):
-    model = models[model_id]
-    processor = processors[model_id]
+# Function to generate prompt
+def generate_caption(image):
+    model = models["microsoft/Phi-3.5-vision-instruct"]
+    processor = processors["microsoft/Phi-3.5-vision-instruct"]
 
-    # Use the default query directly without a user input text field
-    prompt = f"<|user|>\n<|image_1|>\n{DEFAULT_QUERY}<|end|>\n<|assistant|>"
+    prompt = f"<|user|>\n<|image_1|>\n{default_question}<|end|>\n<|assistant|>\n"
     image = Image.fromarray(image).convert("RGB")
 
     inputs = processor(prompt, image, return_tensors="pt").to("cuda:0")
-    generate_ids = model.generate(**inputs, max_new_tokens=1000, eos_token_id=processor.tokenizer.eos_token_id)
+    generate_ids = model.generate(
+        **inputs,
+        max_new_tokens=1000,
+        eos_token_id=processor.tokenizer.eos_token_id,
+    )
     generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
-    response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    response = processor.batch_decode(
+        generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
     return response
 
+# Enhanced CSS for streamlined UI
 css = """
   #container {
     background-color: #f9f9f9;
@@ -48,14 +63,14 @@ css = """
     height: 180px;
     object-fit: contain;
   }
-  #output_text {
+  #output_caption {
     margin-top: 15px;
     border: 2px solid #333;
     border-radius: 8px;
     height: 180px;
     overflow-y: auto;
   }
-  #submit_btn {
+  #run_button {
     background-color: #fff;
     color: black;
     border-radius: 10px;
@@ -64,17 +79,22 @@ css = """
     transition: background-color 0.3s ease;
     margin-top: 15px;
   }
-  #submit_btn:hover {
+  #run_button:hover {
     background-color: #333;
   }
 """
 
+# Gradio Interface with Adjustments
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="container"):
-        input_image = gr.Image(type="pil", label="Upload Image", elem_id="input_image")
-        submit_btn = gr.Button(value="Generate Prompt", elem_id="submit_btn")
-        output_text = gr.Textbox(label="Prompt Output", elem_id="output_text", show_copy_button=True, lines=6)
+        input_image = gr.Image(type="pil", elem_id="input_image", label="Upload Image")
+        run_button = gr.Button(value="Generate Prompt", elem_id="run_button")
+        output_caption = gr.Textbox(label="Generated Prompt", show_copy_button=True, elem_id="output_caption", lines=6)
 
-    submit_btn.click(run_example, [input_image], output_text)
+    run_button.click(
+        fn=generate_caption,
+        inputs=[input_image],
+        outputs=output_caption,
+    )
 
 demo.launch(share=False)