Update app.py

app.py

@@ -1,35 +1,21 @@
-import os
-
-# Install missing dependencies
-try:
-    import torch
-    import torchvision
-    import diffusers
-    import transformers
-    from PIL import Image
-except ModuleNotFoundError as e:
-    missing_module = str(e).split("'")[1]
-    os.system(f"pip install {missing_module}")
-
-import torch
 import streamlit as st
-from diffusers import DiffusionPipeline
-from PIL import Image
-
-# Streamlit UI
-st.title("Lightweight Text-to-Image Generator (CPU Friendly)")
-st.write("Enter a description, and an AI model will generate an image!")
+from diffusers import StableDiffusionPipeline
+import torch
 
 # Load the model
-@st.cache_resource()
+@st.cache_resource
 def load_model():
-    model_id = "CompVis/txt2img-f8-large"  # Lightweight model for CPU
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-    pipe.to("cpu")  # Ensure the model runs on CPU
+    model_id = "CompVis/stable-diffusion-v1-4"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    pipe = pipe.to("cpu")  # Ensure the model runs on CPU
     return pipe
 
 pipe = load_model()
 
+# Streamlit UI
+st.title("Text-to-Image Generator")
+st.write("Enter a description, and the AI model will generate an image!")
+
 # User input
 text_input = st.text_input("Enter your description:")