Update app.py

app.py

@@ -10,16 +10,16 @@ def load_models():
     
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
-    )
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        # device_map="auto"
+    ).to("cuda")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     
     return model, tokenizer
 
 model, tokenizer = load_models()
 
-@spaces.GPU()
+# @spaces.GPU()
 def generate_enhanced_caption(image_caption, text_caption):
     """Generate enhanced caption using the LeX-Enhancer model"""
     combined_caption = f"{image_caption}, with the text on it: {text_caption}."