Update app.py

app.py

@@ -5,23 +5,22 @@ import os
 from PIL import Image
 from huggingface_hub import InferenceClient
 from openai import OpenAI
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
 
-# 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 variable to store the image data URL and prompt for the currently generated image.
 global_image_data_url = None
-global_image_prompt = None
+global_image_prompt = None  # Still stored if needed elsewhere
 
 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.
+    Uses the OpenAI chat model (via Hugging Face Inference API) to generate an image generation prompt
+    based on the selected difficulty, age, and autism 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 Autistic children. "
         f"Consider the following parameters:\n"
@@ -29,15 +28,13 @@ def generate_prompt_from_options(difficulty, age, level):
         f"- Age: {age}\n"
         f"- Autism 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 and helpful."
+        f"can be used to help children learn or understand a concept."
     )
 
     messages = [
         {
             "role": "user",
-            "content": [
-                {"type": "text", "text": query}
-            ]
+            "content": query
         }
     ]
 
@@ -46,7 +43,6 @@ def generate_prompt_from_options(difficulty, age, level):
         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,
@@ -57,7 +53,6 @@ def generate_prompt_from_options(difficulty, age, level):
     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):
@@ -67,22 +62,19 @@ def generate_image_fn(selected_prompt):
     """
     global global_image_data_url, global_image_prompt
 
-    # Save the chosen prompt for later use (for comparison in chat)
+    # Save the chosen prompt for potential future use.
     global_image_prompt = selected_prompt
 
-    # Create an inference client for text-to-image (Stable Diffusion)
     image_client = InferenceClient(
         provider="hf-inference",
         api_key=inference_api_key
     )
 
-    # Generate the image using the selected prompt.
     image = image_client.text_to_image(
         selected_prompt,
         model="stabilityai/stable-diffusion-3.5-large-turbo"
     )
 
-    # Convert the PIL image to a PNG data URL.
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
@@ -93,48 +85,49 @@ def generate_image_fn(selected_prompt):
 
 def generate_image_and_reset_chat(difficulty, age, level, active_session, saved_sessions):
     """
-    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.
+    Saves any current active session into the saved sessions list. Then, using the three selected options,
+    generates an image generation prompt, creates an image, and starts a new active session.
     """
     new_sessions = saved_sessions.copy()
-    # 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):
     """
-    Compares the details entered by the user with the true details (global_image_prompt)
-    and returns hints if needed along with a percentage of correctness.
+    Uses the vision language model to evaluate the user description based solely on the generated image.
+    The message includes both the image (using its data URL) and the user’s text.
     """
-    if not global_image_prompt:
+    if not global_image_data_url:
         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. "
-        "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."
-    )
-
+    # Prepare the message content as a list of parts:
+    # 1. The image part – here we send the image data URL (in practice, you might need to supply a public URL).
+    # 2. The text part – containing the user's description.
     messages = [
         {
             "role": "user",
             "content": [
-                {"type": "text", "text": message_text}
+                {
+                    "type": "image_url",
+                    "image_url": {"url": global_image_data_url}
+                },
+                {
+                    "type": "text",
+                    "text": (
+                        f"Based on the image provided above, please evaluate the following description given by the user:\n"
+                        f"'{user_details}'\n\n"
+                        "Determine a correctness percentage for the description (without referencing the original prompt) "
+                        "and if the description is less than 75% accurate, provide useful hints for improvement."
+                        "Be concise not to overwhelm the user with information."
+                        "you are a kids assistant, so you should be able to explain the image in a simple way."
+                    )
+                }
             ]
         }
     ]
@@ -145,9 +138,9 @@ def compare_details_chat_fn(user_details):
     )
 
     stream = chat_client.chat.completions.create(
-        model="meta-llama/Llama-3.3-70B-Instruct",
+        model="meta-llama/Llama-3.2-11B-Vision-Instruct",
         messages=messages,
-        max_tokens=512,
+        max_tokens=500,
         stream=True
     )
 
@@ -158,9 +151,9 @@ def compare_details_chat_fn(user_details):
 
 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.
+    Processes a new chat message. If no image has been generated yet, instructs the user to generate one.
+    Otherwise, sends the generated image and the user’s description to the vision language model for evaluation,
+    then appends the conversation to the active session's chat history.
     """
     if not active_session.get("image"):
         bot_message = "Please generate an image first."
@@ -184,14 +177,13 @@ def update_sessions(saved_sessions, active_session):
 # Dropdown Options for Difficulty, Age, and Level
 ##############################################
 difficulty_options = ["Easy", "Medium", "Hard"]
-age_options = ["3","4","5","6","7","8","9","10","11","12","13","14","15","16","17","18"]
-level_options = ["Level 1 Autism", "Level 2 Autism", "Level 3 Autism"]
+age_options = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18", "19", "20"]
+level_options = ["Level 1", "Level 2", "Level 3"]
 
 ##############################################
 # Create the Gradio Interface (Single-Page) with a Sidebar for Session Details
 ##############################################
 with gr.Blocks() as demo:
-    # 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([])
 
@@ -219,8 +211,9 @@ 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 description, and the response will indicate "
-                "whether your description is correct, provide hints if needed, and show a percentage of correctness."
+                "Your message will be sent along with the image to a vision language model, "
+                "which will evaluate your description based on what it sees in the image. "
+                "The response will include a correctness percentage and hints if needed."
             )
             chatbot = gr.Chatbot(label="Chat History")
             with gr.Row():