Update app.py

app.py

@@ -2,42 +2,24 @@ import os
 import random
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline  # Changed import
 from transformers import CLIPTextModel, CLIPTokenizer
 
-# Configuration - Using Flux Model
-MODEL_ID = "CompVis/Flux-Pro"
+# Configuration
+MODEL_ID = "CompVis/stable-diffusion-v1-4"  # Changed to working model
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
 def get_pipeline():
-    # Load Flux model components
-    text_encoder = CLIPTextModel.from_pretrained(
+    pipe = StableDiffusionPipeline.from_pretrained(
         MODEL_ID,
-        subfolder="text_encoder",
-        cache_dir=MODEL_CACHE
-    )
-    
-    tokenizer = CLIPTokenizer.from_pretrained(
-        MODEL_ID,
-        subfolder="tokenizer",
-        cache_dir=MODEL_CACHE
-    )
-    
-    # Create pipeline
-    pipe = DiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        cache_dir=MODEL_CACHE,
         torch_dtype=torch.float32,
-        safety_checker=None
+        cache_dir=MODEL_CACHE,
+        safety_checker=None,
+        use_safetensors=True
     )
-    
-    # CPU optimizations
     pipe = pipe.to("cpu")
     pipe.enable_attention_slicing()
-    
     return pipe
 
 # Load model