Update app.py

app.py

@@ -6,7 +6,7 @@ import json
 import torch
 from diffusers import StableDiffusionPipeline
 
-# Get Groq API key from environment variables (set in Space settings)
+# Get Groq API key from environment variables
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 if not GROQ_API_KEY:
     raise ValueError("Please set GROQ_API_KEY in the Space settings under 'Variables'.")
@@ -14,15 +14,20 @@ if not GROQ_API_KEY:
 # Initialize Groq client
 client = Groq(api_key=GROQ_API_KEY)
 
-# Set up device and load the exact model you specified
+# Set up device and load the model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_id = "sd-legacy/stable-diffusion-v1-5"  # Your exact model ID
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,  # Matches your snippet
-    safety_checker=None,  # Optional: disabled for flexibility
-).to(device)
-pipe.enable_model_cpu_offload()  # Optimize memory for Spaces
+print(f"Device available: {device}, CUDA: {torch.cuda.is_available()}")
+model_id = "runwayml/stable-diffusion-v1-5"  # Updated to a widely supported model
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        safety_checker=None,  # Optional: disabled for flexibility
+    ).to(device)
+    pipe.enable_model_cpu_offload()  # Optimize memory usage
+    print(f"Model loaded successfully on {device}: {model_id}")
+except Exception as e:
+    raise ValueError(f"Failed to load model {model_id}: {str(e)}")
 
 # Function to generate tutor output (lesson, question, feedback)
 def generate_tutor_output(subject, difficulty, student_input):
@@ -51,21 +56,28 @@ def generate_tutor_output(subject, difficulty, student_input):
     )
     return completion.choices[0].message.content
 
-# Function to generate images using your model ID
+# Function to generate images
 def generate_images(text, selected_model):
+    print(f"Generating images for text: {text}, model: {selected_model}")
     if selected_model == "Stable Diffusion (Realism)":
         prompt_prefix = "realistic, detailed, vivid colors, "
     elif selected_model == "Stable Diffusion (Portraits)":
         prompt_prefix = "portrait, photorealistic, cinematic lighting, "
     else:
+        print("Invalid model selection")
         return ["Invalid model selection."] * 3
 
     results = []
-    for i in range(3):
-        modified_text = f"{prompt_prefix}{text} variation {i+1}, high quality"
-        image = pipe(modified_text, num_inference_steps=25).images[0]  # Using your model
-        results.append(image)
-    return results
+    try:
+        for i in range(3):
+            modified_text = f"{prompt_prefix}{text} variation {i+1}, high quality"
+            print(f"Generating image {i+1} with prompt: {modified_text}")
+            image = pipe(modified_text, num_inference_steps=25).images[0]
+            results.append(image)
+        return results
+    except Exception as e:
+        print(f"Error in image generation: {str(e)}")
+        return [f"Error: {str(e)}"] * 3
 
 # Gradio interface
 with gr.Blocks(title="AI Tutor with Visuals") as demo:
@@ -112,7 +124,7 @@ 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'.
-    2. **Visual Section**: Select a style and click 'Generate Visuals' to see 3 images from sd-legacy/stable-diffusion-v1-5.
+    2. **Visual Section**: Select a style and click 'Generate Visuals' to see 3 images.
     3. Review the AI-generated content!
 
     *Example*: Try "a photo of an astronaut riding a horse on mars" for visuals.
@@ -131,7 +143,8 @@ with gr.Blocks(title="AI Tutor with Visuals") as demo:
             images = generate_images(text, selected_model)
             return images[0], images[1], images[2]
         except Exception as e:
-            return None, None, None
+            error_msg = f"Visual generation failed: {str(e)}"
+            return error_msg, error_msg, error_msg
 
     submit_button_text.click(
         fn=process_output_text,