Update app.py

app.py

@@ -6,26 +6,31 @@ import json
 import torch
 from diffusers import AutoPipelineForText2Image
 
-# Get API keys from environment variables (set in Space settings)
+# Get Groq API key from environment variables (set in Space settings)
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-# Check if keys are provided
-if not GROQ_API_KEY or not HF_TOKEN:
-    raise ValueError("Please set GROQ_API_KEY and HF_TOKEN in the Space settings under 'Variables'.")
+if not GROQ_API_KEY:
+    raise ValueError("Please set GROQ_API_KEY in the Space settings under 'Variables'.")
 
 # Initialize Groq client
 client = Groq(api_key=GROQ_API_KEY)
 
-# Set up device and image generation pipeline
+# Set up device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipeline = AutoPipelineForText2Image.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.bfloat16,
-    use_auth_token=HF_TOKEN
+
+# 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)
-pipeline.load_lora_weights("Purz/face-projection", weight_name="purz-f4c3_p40j3ct10n.safetensors")
-pipeline.enable_model_cpu_offload()  # Optimize memory usage
+realism_pipeline.enable_model_cpu_offload()  # Optimize memory
+
+photo_pipeline = AutoPipelineForText2Image.from_pretrained(
+    "dreamlike-art/dreamlike-photoreal-2.0",
+    torch_dtype=torch.float16,
+    safety_checker=None,  # Disable NSFW filter (optional)
+).to(device)
+photo_pipeline.enable_model_cpu_offload()  # Optimize memory
 
 # Function to generate tutor output (lesson, question, feedback)
 def generate_tutor_output(subject, difficulty, student_input):
@@ -56,17 +61,19 @@ def generate_tutor_output(subject, difficulty, student_input):
 
 # Function to generate images
 def generate_images(text, selected_model):
-    if selected_model == "Model 1 (Turbo Realism)":
-        prompt_prefix = "realistic, high detail, "
-    elif selected_model == "Model 2 (Face Projection)":
-        prompt_prefix = "f4c3_p40j3ct10n, projection on a face, "
+    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, "
     else:
         return ["Invalid model selection."] * 3
 
     results = []
     for i in range(3):
-        modified_text = f"{prompt_prefix}{text} variation {i+1}"
-        image = pipeline(modified_text, num_inference_steps=20).images[0]
+        modified_text = f"{prompt_prefix}{text} variation {i+1}, high quality"
+        image = pipeline(modified_text, num_inference_steps=25).images[0]
         results.append(image)
     return results
 
@@ -101,9 +108,9 @@ with gr.Blocks(title="AI Tutor with Visuals") as demo:
     with gr.Row():
         with gr.Column(scale=2):
             model_selector = gr.Radio(
-                ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
+                ["Stable Diffusion (Realism)", "Dreamlike Photoreal (Portraits)"],
                 label="Select Image Generation Model",
-                value="Model 1 (Turbo Realism)"
+                value="Stable Diffusion (Realism)"
             )
             submit_button_visual = gr.Button("Generate Visuals", variant="primary")
 
@@ -115,8 +122,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 image variations based on your input.
+    2. **Visual Section**: Select a model and click 'Generate Visuals' to see 3 free, open-source image variations based on your input.
     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.
     """)
 
     def process_output_text(subject, difficulty, student_input):