Update app.py

app.py

@@ -4,30 +4,19 @@ import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
 
-# Install necessary dependencies (for local testing; skip if already installed)
-subprocess.run(
-    "pip install flash-attn --no-build-isolation --global-option='--skip-cuda-build'",
-    shell=True
-)
+# Install flash-attn library
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # Initialize Florence model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-florence_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-base", trust_remote_code=True
-).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained(
-    "microsoft/Florence-2-base", trust_remote_code=True
-)
+florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
+florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
 
-# Define the caption generation function
 def generate_caption(image):
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
-    
-    inputs = florence_processor(
-        text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt"
-    ).to(device)
 
+    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
     generated_ids = florence_model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
@@ -46,48 +35,10 @@ def generate_caption(image):
     print("\n\nGeneration completed!:" + prompt)
     return prompt
 
-# Gradio Interface
-def save_to_csv(images, captions):
-    import csv
-    from io import StringIO
-
-    # Create CSV content
-    output = StringIO()
-    writer = csv.writer(output)
-    writer.writerow(["Filename", "Title", "Keywords"])
-
-    for img, caption in zip(images, captions):
-        filename = img.name if hasattr(img, "name") else "uploaded_image"
-        title = caption[:50]
-        keywords = caption.split(" ")  # Simple keyword generation (replace with a better method)
-        writer.writerow([filename, title, ", ".join(keywords)])
-
-    output.seek(0)
-    return output
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            input_images = gr.Image(
-                label="Upload Images", type="pil", multiple=True
-            )
-            generate_button = gr.Button("Generate Captions")
-        with gr.Column():
-            output_texts = gr.Textbox(
-                label="Generated Captions", lines=5, interactive=False
-            )
-            csv_output = gr.File(label="Download CSV")
-
-    # Define event logic
-    def process(images):
-        captions = [generate_caption(img) for img in images]
-        csv_file = save_to_csv(images, captions)
-        return captions, csv_file
-
-    generate_button.click(
-        fn=process,
-        inputs=[input_images],
-        outputs=[output_texts, csv_output]
-    )
-
-demo.launch(debug=True)
+# Gradio interface
+io = gr.Interface(
+    generate_caption,
+    inputs=[gr.Image(label="Input Image")],
+    outputs=[gr.Textbox(label="Output Prompt", lines=2, show_copy_button=True)]
+)
+io.launch(debug=True)