Update app.py

app.py

@@ -16,8 +16,8 @@ def load_models():
     
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype="auto",
-        device_map="auto"
+        torch_dtype=torch.bfloat16,
+        # device_map="auto"
     )
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     
@@ -26,7 +26,7 @@ def load_models():
         torch_dtype=torch.bfloat16
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe.to("cuda")
+    # pipe.to("cuda")
     
     return model, tokenizer, pipe
 
@@ -43,6 +43,7 @@ def truncate_caption_by_tokens(caption, max_tokens=256):
 
 @spaces.GPU(duration=50)
 def generate_enhanced_caption(image_caption, text_caption):
+    model.to("cuda")
     """Generate enhanced caption using the LeX-Enhancer model"""
     combined_caption = f"{image_caption}, with the text on it: {text_caption}."
     instruction = """
@@ -76,6 +77,7 @@ Below is the simple caption of an image with text. Please deduce the detailed de
 
 @spaces.GPU(duration=60)
 def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
+    pipe.to("cuda")
     """Generate image using LeX-Lumina"""
     # Truncate the caption if it's too long
     enhanced_caption = truncate_caption_by_tokens(enhanced_caption, max_tokens=256)