Update app.py

app.py

@@ -1,17 +1,31 @@
+# app.py
 import gradio as gr
 from groq import Groq
 import os
-import threading  # Import threading module
+import json
+import torch
+from diffusers import AutoPipelineForText2Image
 
-# Initialize Groq client with your API key
-client = Groq(api_key=os.environ["GROQ_API_KEY"])
+# Get API keys from environment variables (set in Space settings)
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Load Text-to-Image Models
-model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
-model2 = gr.load("models/Purz/face-projection")
+# 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'.")
 
-# Stop event for threading (image generation)
-stop_event = threading.Event()
+# Initialize Groq client
+client = Groq(api_key=GROQ_API_KEY)
+
+# Set up device and image generation pipeline
+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  # Authenticate directly with HF token
+).to(device)
+pipeline.load_lora_weights("Purz/face-projection", weight_name="purz-f4c3_p40j3ct10n.safetensors")
+pipeline.enable_model_cpu_offload()  # Optimize memory usage
 
 # Function to generate tutor output (lesson, question, feedback)
 def generate_tutor_output(subject, difficulty, student_input):
@@ -35,117 +49,102 @@ def generate_tutor_output(subject, difficulty, student_input):
             "role": "user",
             "content": prompt,
         }],
-        model="mixtral-8x7b-32768",  # Model for text generation
+        model="mixtral-8x7b-32768",
         max_tokens=1000,
     )
-    
     return completion.choices[0].message.content
 
-# Function to generate images based on model selection
+# Function to generate images
 def generate_images(text, selected_model):
-    stop_event.clear()
-
     if selected_model == "Model 1 (Turbo Realism)":
-        model = model1
+        prompt_prefix = "realistic, high detail, "
     elif selected_model == "Model 2 (Face Projection)":
-        model = model2
+        prompt_prefix = "f4c3_p40j3ct10n, projection on a face, "
     else:
         return ["Invalid model selection."] * 3
 
     results = []
     for i in range(3):
-        if stop_event.is_set():
-            return ["Image generation stopped by user."] * 3
-
-        modified_text = f"{text} variation {i+1}"
-        result = model(modified_text)
-        results.append(result)
-
+        modified_text = f"{prompt_prefix}{text} variation {i+1}"
+        image = pipeline(modified_text, num_inference_steps=20).images[0]  # Faster inference
+        results.append(image)
     return results
 
-# Set up the Gradio interface
-with gr.Blocks() as demo:
+# Gradio interface
+with gr.Blocks(title="AI Tutor with Visuals") as demo:
     gr.Markdown("# 🎓 Your AI Tutor with Visuals & Images")
 
-    # Section for generating Text-based output (lesson, question, feedback)
     with gr.Row():
         with gr.Column(scale=2):
-            # Input fields for subject, difficulty, and student input for textual output
             subject = gr.Dropdown(
-                ["Math", "Science", "History", "Literature", "Code", "AI"], 
-                label="Subject", 
+                ["Math", "Science", "History", "Literature", "Code", "AI"],
+                label="Subject",
                 info="Choose the subject of your lesson"
             )
             difficulty = gr.Radio(
-                ["Beginner", "Intermediate", "Advanced"], 
-                label="Difficulty Level", 
+                ["Beginner", "Intermediate", "Advanced"],
+                label="Difficulty Level",
                 info="Select your proficiency level"
             )
             student_input = gr.Textbox(
-                placeholder="Type your query here...", 
-                label="Your Input", 
+                placeholder="Type your query here...",
+                label="Your Input",
                 info="Enter the topic you want to learn"
             )
             submit_button_text = gr.Button("Generate Lesson & Question", variant="primary")
-        
+
         with gr.Column(scale=3):
-            # Output fields for lesson, question, and feedback
             lesson_output = gr.Markdown(label="Lesson")
             question_output = gr.Markdown(label="Comprehension Question")
             feedback_output = gr.Markdown(label="Feedback")
-    
-    # Section for generating Visual output
+
     with gr.Row():
         with gr.Column(scale=2):
-            # Input fields for text and model selection for image generation
             model_selector = gr.Radio(
                 ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
                 label="Select Image Generation Model",
                 value="Model 1 (Turbo Realism)"
             )
             submit_button_visual = gr.Button("Generate Visuals", variant="primary")
-        
+
         with gr.Column(scale=3):
-            # Output fields for generated images
             output1 = gr.Image(label="Generated Image 1")
             output2 = gr.Image(label="Generated Image 2")
             output3 = gr.Image(label="Generated Image 3")
-    
+
     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, comprehension question, and feedback.
-    2. **Visual Section**: Select the model for image generation, then click 'Generate Visuals' to receive 3 variations of an image based on your topic.
-    3. Review the AI-generated content to enhance your learning experience!
+    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.
+    3. Review the AI-generated content to enhance your learning!
     """)
-    
+
     def process_output_text(subject, difficulty, student_input):
         try:
             tutor_output = generate_tutor_output(subject, difficulty, student_input)
-            parsed = eval(tutor_output)  # Convert string to dictionary
+            parsed = json.loads(tutor_output)
             return parsed["lesson"], parsed["question"], parsed["feedback"]
-        except:
-            return "Error parsing output", "No question available", "No feedback available"
-    
+        except Exception as e:
+            return f"Error: {str(e)}", "No question available", "No feedback available"
+
     def process_output_visual(text, selected_model):
         try:
-            images = generate_images(text, selected_model)  # Generate images
+            images = generate_images(text, selected_model)
             return images[0], images[1], images[2]
-        except:
+        except Exception as e:
             return None, None, None
-    
-    # Generate Text-based Output
+
     submit_button_text.click(
         fn=process_output_text,
         inputs=[subject, difficulty, student_input],
         outputs=[lesson_output, question_output, feedback_output]
     )
-    
-    # Generate Visual Output
+
     submit_button_visual.click(
         fn=process_output_visual,
         inputs=[student_input, model_selector],
         outputs=[output1, output2, output3]
     )
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+# In Hugging Face Spaces, this variable is automatically used as the app entry point
+demo.launch()