Update app.py

app.py

@@ -17,7 +17,7 @@ def load_models():
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.bfloat16,
-        # device_map="auto"
+        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,7 +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")
+    # model.to("cuda")
     """Generate enhanced caption using the LeX-Enhancer model"""
     combined_caption = f"{image_caption}, with the text on it: {text_caption}."
     instruction = """
@@ -77,7 +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")
+    # pipe.to("cuda")
     pipe.enable_model_cpu_offload()
     """Generate image using LeX-Lumina"""
     # Truncate the caption if it's too long