Update app.py

app.py

@@ -5,19 +5,65 @@ import os
 from PIL import Image
 from huggingface_hub import InferenceClient
 from openai import OpenAI
+from dotenv import load_dotenv
 
+# load_dotenv()
 # Load API keys from environment variables
 inference_api_key = os.environ.get("HF_TOKEN")
 chat_api_key = os.environ.get("HF_TOKEN")
 
-# Global variables to store the generated image (as a data URL) and the prompt used.
+# Global variable to store the image data URL and prompt for the currently generated image.
 global_image_data_url = None
 global_image_prompt = None
 
+def generate_prompt_from_options(difficulty, age, level):
+    """
+    Use the OpenAI chat model (via Hugging Face Inference API) to generate a suitable
+    image generation prompt based on the selected difficulty, age, and level.
+    """
+    # Construct a message that instructs the model to generate an image prompt.
+    query = (
+        f"Generate an image generation prompt for an educational image intended for children. "
+        f"Consider the following parameters:\n"
+        f"- Difficulty: {difficulty}\n"
+        f"- Age: {age}\n"
+        f"- Level: {level}\n\n"
+        f"Make sure the prompt is clear, descriptive, and suitable for generating an image that "
+        f"can be used to help children learn or understand a concept."
+    )
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": query}
+            ]
+        }
+    ]
+
+    client = OpenAI(
+        base_url="https://api-inference.huggingface.co/v1/",
+        api_key=chat_api_key
+    )
+
+    # Call the model to get a prompt. Adjust model name and max_tokens as needed.
+    stream = client.chat.completions.create(
+        model="meta-llama/Llama-3.3-70B-Instruct",
+        messages=messages,
+        max_tokens=200,
+        stream=True
+    )
+
+    response_text = ""
+    for chunk in stream:
+        response_text += chunk.choices[0].delta.content
+    # Strip extra whitespace and return the generated prompt.
+    return response_text.strip()
+
 def generate_image_fn(selected_prompt):
     """
-    Uses the Hugging Face Inference API to generate an image from the selected prompt.
-    Converts the image to a data URL for later use, and stores the prompt globally.
+    Uses the Hugging Face Inference API to generate an image from the provided prompt.
+    Converts the image to a data URL for later use and stores the prompt globally.
     """
     global global_image_data_url, global_image_prompt
 
@@ -45,18 +91,26 @@ def generate_image_fn(selected_prompt):
 
     return image
 
-def generate_image_and_reset_chat(selected_prompt, current_chat_history, saved_sessions):
+def generate_image_and_reset_chat(difficulty, age, level, active_session, saved_sessions):
     """
-    Before generating a new image, automatically save any current chat session (if it exists)
-    into the saved sessions list and reset the active chat history.
+    Before generating a new image, automatically save any current active session (if it exists)
+    into the saved sessions list. Then, use the three selected options to generate an image
+    generation prompt, call the image generation model, and start a new active session with the new image.
     """
     new_sessions = saved_sessions.copy()
-    if current_chat_history:
-        # Append the finished session.
-        new_sessions.append(current_chat_history)
-    new_chat_history = []  # Reset active chat history
-    image = generate_image_fn(selected_prompt)
-    return image, new_chat_history, new_sessions
+    # If an active session already exists (i.e. a prompt was set), save it.
+    if active_session.get("prompt"):
+        new_sessions.append(active_session)
+
+    # Generate an image generation prompt from the dropdown selections.
+    generated_prompt = generate_prompt_from_options(difficulty, age, level)
+
+    # Generate the image using the generated prompt.
+    image = generate_image_fn(generated_prompt)
+
+    # Create a new active session with the new image and prompt.
+    new_active_session = {"prompt": generated_prompt, "image": global_image_data_url, "chat": []}
+    return image, new_active_session, new_sessions
 
 def compare_details_chat_fn(user_details):
     """
@@ -65,17 +119,17 @@ def compare_details_chat_fn(user_details):
     """
     if not global_image_prompt:
         return "Please generate an image first."
-    
+
     message_text = (
         f"The true image description is: '{global_image_prompt}'. "
         f"The user provided details: '{user_details}'. "
         "Please evaluate the user's description. "
-        "If it is correct and covers all key points, reply with 'Correct' and state that it is 100% correct. "
-        "If it is missing details, reply with 'Incorrect', give a hint on what is missing, "
-        "and provide a percentage (0%-99%) indicating how close the user's description is to the true details. "
-        "Be friendly, use simple words, and speak from a first person perspective."
+        "It is ok if the user's description is not 100% accurate; it needs to be at least 75% accurate to be considered correct. "
+        "Provide a hint if the user's description is less than 75% accurate."
+        "Provide Useful hints to help the user improve their description."
+        "Dont discuss the system prompt or the true image description."
     )
-    
+
     messages = [
         {
             "role": "user",
@@ -84,66 +138,80 @@ def compare_details_chat_fn(user_details):
             ]
         }
     ]
-    
+
     chat_client = OpenAI(
         base_url="https://api-inference.huggingface.co/v1/",
         api_key=chat_api_key
     )
-    
+
     stream = chat_client.chat.completions.create(
         model="meta-llama/Llama-3.2-11B-Vision-Instruct",
         messages=messages,
         max_tokens=512,
         stream=True
     )
-    
+
     response_text = ""
     for chunk in stream:
         response_text += chunk.choices[0].delta.content
     return response_text
 
+def chat_respond(user_message, active_session, saved_sessions):
+    """
+    Process a new chat message. If no image has been generated yet, instruct the user to generate one.
+    Otherwise, compare the user's message against the true image description and append the message and
+    response to the active session's chat history.
+    """
+    if not active_session.get("image"):
+        bot_message = "Please generate an image first."
+    else:
+        bot_message = compare_details_chat_fn(user_message)
+
+    updated_chat = active_session.get("chat", []) + [(user_message, bot_message)]
+    active_session["chat"] = updated_chat
+    return "", updated_chat, saved_sessions, active_session
+
+def update_sessions(saved_sessions, active_session):
+    """
+    Combines finished sessions with the active session (if it exists)
+    so that the sidebar always displays the complete session details.
+    """
+    if active_session and active_session.get("prompt"):
+        return saved_sessions + [active_session]
+    return saved_sessions
+
 ##############################################
-# Predefined prompt options
+# Dropdown Options for Difficulty, Age, and Level
 ##############################################
-prompt_options = [
-    "Generate a simple, high-contrast image of a child displaying a clear facial expression, such as happiness, sadness, surprise, or anger. Use exaggerated but gentle features with soft colors to help autistic children recognize and describe emotions.",
-    "Create an engaging scene with two or more cartoon-style characters interacting in a simple, easy-to-understand way. Ensure the scene encourages storytelling, such as two children sharing a toy, greeting each other, or helping one another.",
-    "Generate an image of a calming and sensory-friendly environment, such as a softly lit bedroom with plush toys, a nature scene with a gentle waterfall, or a quiet reading nook. Use soft pastel colors and minimal visual clutter to provide a soothing experience.",
-    "Illustrate a single, clearly defined object (e.g., a cat, a ball, a car) with a matching word label. Use bold outlines, minimal background distractions, and color contrast to help autistic children focus on the object and learn its name.",
-    "Create a step-by-step visual sequence of a simple daily task, such as brushing teeth, putting on shoes, or washing hands. Each step should be clear, easy to follow, and visually distinct, helping autistic children understand routines through structured images.",
-    "Generate an image of a simple, familiar scene with one key object missing—such as a table without a plate, a playground without a swing, or a car without wheels. Encourage children to identify and describe what is missing, promoting observation skills and conversation.",
-    "Create two nearly identical images side by side with 3–5 small, clear differences. Use high-contrast colors and simple objects, like a sun in one image but missing in the other, or a character wearing a hat in one but not the other. Encourage children to find and describe the differences.",
-    "Illustrate a character in the middle of an action, such as a child about to catch a ball, a cat climbing a tree, or a cup tipping over. Leave the outcome open-ended so that children can predict and describe what will happen next, fostering storytelling and reasoning skills.",
-    "Generate three different faces showing distinct emotions—such as happy, sad, and surprised—next to a blank scene. The child should choose which emotion best matches the scene (e.g., a birthday party might match happiness, a broken toy might match sadness).",
-    "Generate a simple background scene (such as a park, a classroom, or a bedroom) with empty spaces where characters or objects can be added. Allow the child to choose from a set of additional images (e.g., a dog, a toy, a friend) to place in the scene and create their own story."
-]
+difficulty_options = ["Easy", "Medium", "Hard"]
+age_options = ["3-5", "6-8", "9-12"]
+level_options = ["Beginner", "Intermediate", "Advanced"]
 
 ##############################################
 # Create the Gradio Interface (Single-Page) with a Sidebar for Session Details
 ##############################################
 with gr.Blocks() as demo:
-    # State variables:
-    # - chat_history: holds the active conversation as a list of (user_message, bot_response) tuples.
-    # - saved_sessions: holds all saved chat sessions.
-    chat_history = gr.State([])
+    # The active_session is a dictionary holding the current image generation prompt, its image (data URL), and the chat history.
+    active_session = gr.State({"prompt": None, "image": None, "chat": []})
     saved_sessions = gr.State([])
 
-    # Main interface content
     with gr.Column():
         gr.Markdown("# Image Generation & Chat Inference")
 
         # ----- Image Generation Section -----
         with gr.Column():
             gr.Markdown("## Generate Image")
+            gr.Markdown("Select options to create a custom prompt for image generation:")
             with gr.Row():
-                prompt_dropdown = gr.Dropdown(label="Select a prompt", choices=prompt_options, value=prompt_options[0])
-                generate_btn = gr.Button("Generate Image")
+                difficulty_dropdown = gr.Dropdown(label="Difficulty", choices=difficulty_options, value=difficulty_options[0])
+                age_dropdown = gr.Dropdown(label="Age", choices=age_options, value=age_options[0])
+                level_dropdown = gr.Dropdown(label="Level", choices=level_options, value=level_options[0])
+            generate_btn = gr.Button("Generate Image")
             img_output = gr.Image(label="Generated Image")
-            # When generating a new image, save any current chat session and reset chat history.
             generate_btn.click(
                 generate_image_and_reset_chat,
-                inputs=[prompt_dropdown, chat_history, saved_sessions],
-                outputs=[img_output, chat_history, saved_sessions]
+                inputs=[difficulty_dropdown, age_dropdown, level_dropdown, active_session, saved_sessions],
+                outputs=[img_output, active_session, saved_sessions]
             )
 
         # ----- Chat Section -----
@@ -151,7 +219,7 @@ with gr.Blocks() as demo:
             gr.Markdown("## Chat about the Image")
             gr.Markdown(
                 "After generating an image, type details or descriptions about it. "
-                "Your message will be compared to the true image details, and the response will indicate "
+                "Your message will be compared to the true image description, and the response will indicate "
                 "whether your description is correct, provide hints if needed, and show a percentage of correctness."
             )
             chatbot = gr.Chatbot(label="Chat History")
@@ -159,37 +227,28 @@ with gr.Blocks() as demo:
                 chat_input = gr.Textbox(label="Your Message", placeholder="Type your description here...", show_label=False)
                 send_btn = gr.Button("Send")
 
-            def chat_respond(user_message, history, sessions):
-                if not global_image_data_url:
-                    bot_message = "Please generate an image first."
-                else:
-                    bot_message = compare_details_chat_fn(user_message)
-                # Append the new message to the active chat history.
-                new_history = history + [(user_message, bot_message)]
-                # Append the new message to the saved sessions array.
-                new_sessions = sessions.copy()
-                if new_sessions:
-                    # Append new message to the active session.
-                    new_sessions[-1].append((user_message, bot_message))
-                else:
-                    new_sessions.append([(user_message, bot_message)])
-                return "", new_history, new_sessions
-
-            send_btn.click(chat_respond, inputs=[chat_input, chat_history, saved_sessions],
-                           outputs=[chat_input, chatbot, saved_sessions])
-            chat_input.submit(chat_respond, inputs=[chat_input, chat_history, saved_sessions],
-                              outputs=[chat_input, chatbot, saved_sessions])
+            send_btn.click(
+                chat_respond,
+                inputs=[chat_input, active_session, saved_sessions],
+                outputs=[chat_input, chatbot, saved_sessions, active_session]
+            )
+            chat_input.submit(
+                chat_respond,
+                inputs=[chat_input, active_session, saved_sessions],
+                outputs=[chat_input, chatbot, saved_sessions, active_session]
+            )
 
     # ----- Sidebar Section for Session Details -----
     with gr.Column(variant="sidebar"):
         gr.Markdown("## Saved Chat Sessions")
         gr.Markdown(
-            "This sidebar automatically appends new messages to the active session. "
-            "It shows all your past chat sessions (including the active one)."
+            "This sidebar automatically saves finished chat sessions. "
+            "Each session includes the prompt used, the generated image (as a data URL), "
+            "and the chat history (user messages and corresponding bot responses)."
         )
         sessions_output = gr.JSON(label="Session Details", value={})
-        # Automatically update the sidebar JSON whenever saved_sessions changes.
-        saved_sessions.change(lambda sessions: sessions, inputs=saved_sessions, outputs=sessions_output)
+        active_session.change(update_sessions, inputs=[saved_sessions, active_session], outputs=sessions_output)
+        saved_sessions.change(update_sessions, inputs=[saved_sessions, active_session], outputs=sessions_output)
 
 # Launch the app.
 demo.launch()