Update app.py

app.py

@@ -4,7 +4,7 @@ from groq import Groq
 import os
 import json
 import torch
-from diffusers import AutoPipelineForText2Image
+from diffusers import StableDiffusionPipeline
 
 # Get Groq API key from environment variables (set in Space settings)
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
@@ -14,23 +14,14 @@ if not GROQ_API_KEY:
 # Initialize Groq client
 client = Groq(api_key=GROQ_API_KEY)
 
-# Set up device
+# Set up device and load Stable Diffusion model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load two open-access models
-realism_pipeline = AutoPipelineForText2Image.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16,  # Faster on GPU
-    safety_checker=None,  # Disable NSFW filter (optional; comment out if unwanted)
-).to(device)
-realism_pipeline.enable_model_cpu_offload()  # Optimize memory
-
-photo_pipeline = AutoPipelineForText2Image.from_pretrained(
-    "dreamlike-art/dreamlike-photoreal-2.0",
+pipe = StableDiffusionPipeline.from_pretrained(
+    "sd-legacy/stable-diffusion-v1-5",
     torch_dtype=torch.float16,
-    safety_checker=None,  # Disable NSFW filter (optional)
+    safety_checker=None,  # Optional: disable NSFW filter (comment out if unwanted)
 ).to(device)
-photo_pipeline.enable_model_cpu_offload()  # Optimize memory
+pipe.enable_model_cpu_offload()  # Optimize memory usage
 
 # Function to generate tutor output (lesson, question, feedback)
 def generate_tutor_output(subject, difficulty, student_input):
@@ -59,21 +50,19 @@ def generate_tutor_output(subject, difficulty, student_input):
     )
     return completion.choices[0].message.content
 
-# Function to generate images
+# Function to generate images using sd-legacy/stable-diffusion-v1-5
 def generate_images(text, selected_model):
     if selected_model == "Stable Diffusion (Realism)":
-        pipeline = realism_pipeline
         prompt_prefix = "realistic, detailed, vivid colors, "
-    elif selected_model == "Dreamlike Photoreal (Portraits)":
-        pipeline = photo_pipeline
-        prompt_prefix = "photorealistic portrait, cinematic lighting, "
+    elif selected_model == "Stable Diffusion (Portraits)":
+        prompt_prefix = "portrait, photorealistic, cinematic lighting, "
     else:
         return ["Invalid model selection."] * 3
 
     results = []
     for i in range(3):
         modified_text = f"{prompt_prefix}{text} variation {i+1}, high quality"
-        image = pipeline(modified_text, num_inference_steps=25).images[0]
+        image = pipe(modified_text, num_inference_steps=25).images[0]
         results.append(image)
     return results
 
@@ -108,8 +97,8 @@ with gr.Blocks(title="AI Tutor with Visuals") as demo:
     with gr.Row():
         with gr.Column(scale=2):
             model_selector = gr.Radio(
-                ["Stable Diffusion (Realism)", "Dreamlike Photoreal (Portraits)"],
-                label="Select Image Generation Model",
+                ["Stable Diffusion (Realism)", "Stable Diffusion (Portraits)"],
+                label="Select Image Style",
                 value="Stable Diffusion (Realism)"
             )
             submit_button_visual = gr.Button("Generate Visuals", variant="primary")
@@ -122,10 +111,10 @@ with gr.Blocks(title="AI Tutor with Visuals") as demo:
     gr.Markdown("""
     ### How to Use
     1. **Text Section**: Select a subject and difficulty, type your query, and click 'Generate Lesson & Question' to get your personalized lesson, question, and feedback.
-    2. **Visual Section**: Select a model and click 'Generate Visuals' to see 3 free, open-source image variations based on your input.
+    2. **Visual Section**: Select an image style and click 'Generate Visuals' to see 3 variations using Stable Diffusion v1.5 (free & open-source).
     3. Review the AI-generated content to enhance your learning!
 
-    *Note*: These use free, open-access models (Stable Diffusion & Dreamlike Photoreal). GPU recommended for speed.
+    *Note*: Powered by sd-legacy/stable-diffusion-v1-5. GPU recommended for faster results.
     """)
 
     def process_output_text(subject, difficulty, student_input):