Update app.py

app.py

@@ -19,8 +19,8 @@ def load_text_model():
         # Load model with FP16 or 8-bit quantization
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,  # Reduce VRAM usage
-            low_cpu_mem_usage=True  # Optimize memory
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            low_cpu_mem_usage=True 
         ).to(device)
 
         st.write("✅ Text model loaded successfully!")
@@ -40,7 +40,7 @@ def load_image_model():
         model_id = "runwayml/stable-diffusion-v1-5"
         model = StableDiffusionPipeline.from_pretrained(
             model_id,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32  # Reduce VRAM usage
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
         ).to(device)
         model.enable_attention_slicing()  # Optimize GPU memory
         st.write("✅ Image model loaded successfully!")