Update app.py

app.py

@@ -186,4 +186,44 @@ with gr.Blocks(theme=theme, css=css, elem_id="app-container") as app:
         
         text_button.click(query, inputs=[custom_lora, text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=[image_output, seed_output])
 
+    with gr.Tab("Image to Prompt"): 
+            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)
+    
+            # api_key = os.getenv("HF_READ_TOKEN")
+            
+            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)
+                generated_ids = florence_model.generate(
+                    input_ids=inputs["input_ids"],
+                    pixel_values=inputs["pixel_values"],
+                    max_new_tokens=1024,
+                    early_stopping=False,
+                    do_sample=False,
+                    num_beams=3,
+                )
+                generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+                parsed_answer = florence_processor.post_process_generation(
+                    generated_text,
+                    task="<MORE_DETAILED_CAPTION>",
+                    image_size=(image.width, image.height)
+                )
+                prompt =  parsed_answer["<MORE_DETAILED_CAPTION>"]
+                print("\n\nGeneration completed!:"+ prompt)
+                return prompt
+    
+            io = gr.Interface(generate_caption,
+                inputs=[gr.Image(label="Input Image")],
+                outputs = [gr.Textbox(label="Output Prompt", lines=2, show_copy_button = True),
+                # gr.Image(label="Output Image")
+                           ]
+                        )
+
 app.launch(show_api=False, share=False)