Update app.py

app.py

@@ -6,15 +6,24 @@ import io
 from PIL import Image
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-#
-model_id = 'J-LAB/Florence-vl3'
-model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+model_ids = {
+    "Florence-vl2": 'J-LAB/Florence-vl2',
+    "Florence-vl3": 'J-LAB/Florence-vl3',
+    "Florence_2_F_FluxiAI_Product_Caption": 'J-LAB/Florence_2_F_FluxiAI_Product_Caption'
+}
+
+# Load model and processor based on the selected model
+def load_model(model_name):
+    model_id = model_ids[model_name]
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
+    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    return model, processor
 
 DESCRIPTION = "# Product Describe by Fluxi IA\n### Base Model [Florence-2] (https://huggingface.co/microsoft/Florence-2-large)"
 
 @spaces.GPU
-def run_example(task_prompt, image):
+def run_example(model, processor, task_prompt, image):
     inputs = processor(text=task_prompt, images=image, return_tensors="pt").to("cuda")
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
@@ -32,14 +41,16 @@ def run_example(task_prompt, image):
     )
     return parsed_answer
 
-def process_image(image, task_prompt):
+def process_image(image, task_prompt, model_name):
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    model, processor = load_model(model_name)
+    
     if task_prompt == 'Product Caption':
         task_prompt = '<MORE_DETAILED_CAPTION>'
     elif task_prompt == 'OCR':
         task_prompt = '<OCR>'
     
-    results = run_example(task_prompt, image)
+    results = run_example(model, processor, task_prompt, image)
     
     # Remove the key and get the text value
     if results and task_prompt in results:
@@ -80,11 +91,16 @@ single_task_list =[
     'Product Caption', 'OCR'
 ]
 
+model_list = [
+    'Florence-vl2', 'Florence-vl3', 'Florence_2_F_FluxiAI_Product_Caption'
+]
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Tab(label="Product Image Select"):
         with gr.Row():
             with gr.Column():
+                model_name = gr.Dropdown(choices=model_list, label="Model", value="Florence-vl3")
                 input_img = gr.Image(label="Input Picture")
                 task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Product Caption")
                 submit_btn = gr.Button(value="Submit")
@@ -108,8 +124,8 @@ with gr.Blocks(css=css) as demo:
         ```
         """)
         
-        submit_btn.click(process_image, [input_img, task_prompt], [output_text])
+        submit_btn.click(process_image, [input_img, task_prompt, model_name], [output_text])
 
     demo.load(lambda: None, inputs=None, outputs=None, js=js)
 
-demo.launch(debug=True) 
+demo.launch(debug=True)