Update app.py

app.py

@@ -2,6 +2,7 @@ import spaces
 import gradio as gr
 import torch
 from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, pipeline
+from transformers import AutoProcessor, AutoModelForCausalLM
 import re
 import random
 import os
@@ -40,12 +41,35 @@ kolors_pipe.enable_model_cpu_offload()
 vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner-v2").to(device).eval()
 vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner-v2")
 
+# Initialize Florence model
+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)
+
 # Prompt Enhancer
 enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
 enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
 
 MAX_SEED = 2**32 - 1
 
+# Florence caption function
+def florence_caption(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)
+    )
+    return parsed_answer["<MORE_DETAILED_CAPTION>"]
+
 # VLM Captioner function
 def create_captions_rich(image):
     prompt = "caption en"
@@ -112,9 +136,12 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
 
 # Gradio Interface
 @spaces.GPU
-def process_workflow(image, text_prompt, use_vlm, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def process_workflow(image, text_prompt, use_vlm, use_enhancer, model_choice, vlm_model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
     if use_vlm and image is not None:
-        prompt = create_captions_rich(image)
+        if vlm_model_choice == "Long Captioner":
+            prompt = create_captions_rich(image)
+        else:  # Florence
+            prompt = florence_caption(image)
     else:
         prompt = text_prompt
     
@@ -161,6 +188,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
             with gr.Group(elem_classes="input-group"):
                 input_image = gr.Image(label="Input Image for VLM")
                 use_vlm = gr.Checkbox(label="Use VLM Captioner", value=False)
+                vlm_model_choice = gr.Radio(["Long Captioner", "Florence"], label="VLM Model", value="Long Captioner")
             
             with gr.Group(elem_classes="input-group"):
                 text_prompt = gr.Textbox(label="Text Prompt")
@@ -187,7 +215,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
     generate_btn.click(
         fn=process_workflow,
         inputs=[
-            input_image, text_prompt, use_vlm, use_enhancer, model_choice,
+            input_image, text_prompt, use_vlm, use_enhancer, model_choice, vlm_model_choice,
             negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps
         ],
         outputs=[output_image, final_prompt, used_seed]