Update app.py

app.py

@@ -24,11 +24,11 @@ def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
     imports.remove("flash_attn")
     return imports
 
-
-models = {
-    'microsoft/Florence-2-base-ft': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True, device_map='cpu'),
-    'microsoft/Florence-2-base': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True, device_map='cpu'),
-}
+with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+    models = {
+        'microsoft/Florence-2-base-ft': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True, device_map='cpu'),
+        'microsoft/Florence-2-base': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True, device_map='cpu'),
+    }
 
 processors = {
     'microsoft/Florence-2-base-ft': AutoProcessor.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True),
@@ -49,14 +49,13 @@ def fig_to_pil(fig):
 
 
 def run_example(task_prompt, image, text_input=None, model_id='microsoft/Florence-2-base'):
-    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
-        model = models[model_id]
-        processor = processors[model_id]
+    model = models[model_id]
+    processor = processors[model_id]
     if text_input is None:
         prompt = task_prompt
     else:
         prompt = task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],