Update app.py

app.py

@@ -22,8 +22,13 @@ import requests
 from io import BytesIO
 import json
 import torch
-from diffusers import DiffusionPipeline
-import torch
+
+# Use compatible import for diffusers
+try:
+    from diffusers import StableDiffusionPipeline
+except ImportError:
+    from diffusers import DiffusionPipeline as StableDiffusionPipeline
+
 from transformers import pipeline
 
 # Configure Streamlit page
@@ -230,13 +235,13 @@ def load_models():
     
     try:
         # Use a fast, optimized model
-        models['text_to_image'] = DiffusionPipeline.from_pretrained(
+        models['text_to_image'] = StableDiffusionPipeline.from_pretrained(
             "OFA-Sys/small-stable-diffusion-v0",  # Optimized for speed
             torch_dtype=torch.float32,
             safety_checker=None
         )
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        models['text_to_image'].to(device)
+        models['text_to_image'] = models['text_to_image'].to(device)
     except Exception as e:
         st.error(f"Could not load text-to-image model: {str(e)}")
         models['text_to_image'] = None