Update app.py

app.py

@@ -25,17 +25,12 @@ model_loaded = False
 load_error_message = ""
 
 try:
-    # We recommend enabling flash_attention_2 for better acceleration and memory saving if available.
-    # For broader compatibility, we'll try without it first, or conditionally enable it.
-    # attn_implementation = "flash_attention_2" if torch.cuda.is_available() and hasattr(torch.nn.functional, 'scaled_dot_product_attention') else None
-    
     model = AutoModelForImageTextToText.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
-        device_map="auto",
         trust_remote_code=True 
-    )
+    ).to("cuda")
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
     model_loaded = True
     print("Model and processor loaded successfully.")
@@ -75,6 +70,7 @@ def run_inference_localization(
     messages_for_template: List[dict[str, Any]], 
     pil_image_for_processing: Image.Image
 ) -> str:
+    torch.cuda.set_device("cuda")
     """
     Runs inference using the Holo1 model.
     - messages_for_template: The prompt structure, potentially including the PIL image object