Update app.py

app.py

@@ -6,24 +6,15 @@ 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_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
+#
+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)
 
 DESCRIPTION = "# Product Describe by Fluxi IA\n### Base Model [Florence-2] (https://huggingface.co/microsoft/Florence-2-large)"
 
 @spaces.GPU
-def run_example(model, processor, task_prompt, image):
+def run_example(task_prompt, image):
     inputs = processor(text=task_prompt, images=image, return_tensors="pt").to("cuda")
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
@@ -41,16 +32,14 @@ def run_example(model, processor, task_prompt, image):
     )
     return parsed_answer
 
-def process_image(image, task_prompt, model_name):
+def process_image(image, task_prompt):
     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(model, processor, task_prompt, image)
+    results = run_example(task_prompt, image)
     
     # Remove the key and get the text value
     if results and task_prompt in results:
@@ -91,16 +80,11 @@ 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")
@@ -124,8 +108,8 @@ with gr.Blocks(css=css) as demo:
         ```
         """)
         
-        submit_btn.click(process_image, [input_img, task_prompt, model_name], [output_text])
+        submit_btn.click(process_image, [input_img, task_prompt], [output_text])
 
     demo.load(lambda: None, inputs=None, outputs=None, js=js)
 
-demo.launch(debug=True)
+demo.launch(debug=True)