Update app.py

app.py

@@ -1,55 +1,51 @@
 import streamlit as st
+from diffusers import StableDiffusionPipeline
 import torch
-from transformers import CLIPProcessor, CLIPTextModel  # Using CLIP as an example of a lighter model for image generation
 
 @st.cache_resource
-def load_lightweight_model():
-    """Load a lightweight image generation model."""
+def load_lightweight_pipeline():
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model_name = "CompVis/ldm-text2im-large-256"  # A known lightweight diffusion model
-    model = CLIPTextModel.from_pretrained(model_name).to(device)
-    processor = CLIPProcessor.from_pretrained(model_name)
-    return model, processor
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4", 
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    ).to(device)
+    return pipeline
 
-def generate_image(model, processor, prompt, negative_prompt):
-    """Generate an image using the lightweight model."""
+def generate_image(pipeline, prompt, negative_prompt):
     try:
-        inputs = processor(text=prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(inputs['input_ids'], max_length=50)
-        # Here we would simulate or replace this with the actual image output from the lightweight model
-        dummy_image = "https://via.placeholder.com/256"  # Placeholder image link for demonstration
-        return dummy_image
+        image = pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=25,
+            guidance_scale=7.5
+        ).images[0]
+        return image
     except Exception as e:
         st.error(f"Error generating image: {str(e)}")
         return None
 
-# Streamlit app setup
 st.title("Image Generator")
-st.markdown("### By Taizun", unsafe_allow_html=True)
-st.markdown('<small>Effortlessly create stunning images with our lightweight generator.</small>', unsafe_allow_html=True)
+st.markdown('<small>By Taizun</small>', unsafe_allow_html=True)
 
 prompt = st.text_input("Enter your prompt")
-model, processor = load_lightweight_model()
+pipeline = load_lightweight_pipeline()
 
-# Style table (Example styles)
 styles_dict = {
-    "Neon Vibes": {"prompt": "neon lights", "negative_prompt": "blurry, low quality"},
-    "Retro Sci-Fi": {"prompt": "retro futuristic scene", "negative_prompt": "modern, dull"},
-    "Mystic Forest": {"prompt": "dark forest with mystical lights", "negative_prompt": "bright, plain"},
-    "Abstract Art": {"prompt": "abstract shapes and colors", "negative_prompt": "realistic, boring"},
+    "Neon Vibes": {"prompt": "neon lights in a futuristic city", "negative_prompt": "blurred, dull"},
+    "Retro Sci-Fi": {"prompt": "retro futuristic landscape", "negative_prompt": "realistic, modern"},
+    "Mystic Forest": {"prompt": "dark mystical forest with glowing lights", "negative_prompt": "sunny, bright"},
+    "Abstract Art": {"prompt": "abstract colorful shapes", "negative_prompt": "realistic, detailed"}
 }
 
-# Dropdown for selecting a style
 style_name = st.selectbox("Select a Style", options=list(styles_dict.keys()))
 
-# Display the selected style's prompt and negative prompt
 if style_name:
-    selected_entry = styles_dict[style_name]
-    selected_style_prompt = selected_entry["prompt"]
-    selected_style_negative_prompt = selected_entry["negative_prompt"]
+    selected_style = styles_dict[style_name]
+    style_prompt = selected_style["prompt"]
+    negative_prompt = selected_style["negative_prompt"]
 
 if st.button("Generate Image"):
     with st.spinner("Generating your image..."):
-        result = generate_image(model, processor, prompt + " " + selected_style_prompt, selected_style_negative_prompt)
-        if result:
-            st.image(result, caption="Generated Image", use_column_width=True)
+        image = generate_image(pipeline, prompt + " " + style_prompt, negative_prompt)
+        if image:
+            st.image(image, caption="Generated Image", use_column_width=True)