Update app.py

app.py

@@ -7,30 +7,42 @@ from diffusers import StableDiffusionPipeline
 device = "cuda" if torch.cuda.is_available() else "cpu"
 st.write(f"Using device: {device}")  # Debug message
 
-# Load text model (TinyLlama) with error handling
+# Load text model (TinyLlama) with optimizations
 @st.cache_resource
 def load_text_model():
     try:
         st.write("⏳ Loading text model...")
         model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+        
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+        
+        # 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
+        ).to(device)
+
         st.write("✅ Text model loaded successfully!")
-        return pipeline("text-generation", model=model, tokenizer=tokenizer)
+        return pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if device == "cuda" else -1)
+    
     except Exception as e:
         st.error(f"❌ Error loading text model: {e}")
         return None
 
 story_generator = load_text_model()
 
-# Load image model (Stable Diffusion) with error handling
+# Load image model (Stable Diffusion) with optimizations
 @st.cache_resource
 def load_image_model():
     try:
         st.write("⏳ Loading image model...")
         model_id = "runwayml/stable-diffusion-v1-5"
-        model = StableDiffusionPipeline.from_pretrained(model_id).to(device)
-        model.enable_attention_slicing()  # Optimize memory usage
+        model = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32  # Reduce VRAM usage
+        ).to(device)
+        model.enable_attention_slicing()  # Optimize GPU memory
         st.write("✅ Image model loaded successfully!")
         return model
     except Exception as e: