Update app.py

app.py

@@ -21,13 +21,13 @@ def generate_image_fn(selected_prompt):
     """
     global global_image_data_url, global_image_prompt
 
-    # Store the chosen prompt for later use in detail checking
+    # Save the chosen prompt for later 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  # Loaded from env secrets
+        api_key=inference_api_key
     )
 
     # Generate the image using the selected prompt.
@@ -36,7 +36,7 @@ def generate_image_fn(selected_prompt):
         model="stabilityai/stable-diffusion-3.5-large-turbo"
     )
 
-    # Convert the PIL image to a PNG data URL
+    # Convert the PIL image to a PNG data URL.
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
@@ -45,16 +45,15 @@ def generate_image_fn(selected_prompt):
 
     return image
 
-def generate_image_and_reset_chat(selected_prompt, current_chat_history, current_chat_sessions):
+def generate_image_and_reset_chat(selected_prompt, current_chat_history, saved_sessions):
     """
-    Before generating a new image, automatically save the current chat session (if any)
-    into the sessions state, then clear the chat history.
-    Returns the generated image along with updated chat history (empty) and chat sessions.
+    Before generating a new image, automatically save any current chat session (if exists)
+    into the saved sessions list and reset the active chat history.
     """
-    new_sessions = current_chat_sessions.copy()
+    new_sessions = saved_sessions.copy()
     if current_chat_history:
         new_sessions.append(current_chat_history)
-    new_chat_history = []
+    new_chat_history = []  # Reset active chat history
     image = generate_image_fn(selected_prompt)
     return image, new_chat_history, new_sessions
 
@@ -81,63 +80,13 @@ def chat_about_image_fn(user_input):
 
     chat_client = OpenAI(
         base_url="https://api-inference.huggingface.co/v1/",
-        api_key=chat_api_key  # Loaded from env secrets
-    )
-
-    stream = chat_client.chat.completions.create(
-        model="meta-llama/Llama-3.2-11B-Vision-Instruct", 
-        messages=messages, 
-        max_tokens=500,
-        stream=True
-    )
-
-    response_text = ""
-    for chunk in stream:
-        response_text += chunk.choices[0].delta.content
-
-    return response_text
-
-def check_details_fn(user_details):
-    """
-    Compares the user's description of the generated image with the prompt used to generate it.
-    The function sends both the original prompt and the user description to the vision-chat model,
-    which responds whether the description is correct and (if not) provides a hint.
-    """
-    if not global_image_prompt:
-        return "Please generate an image first."
-
-    # Build a message to instruct the model to evaluate the user's details.
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": (
-                        f"The image was generated using the prompt: '{global_image_prompt}'.\n"
-                        f"Evaluate the following user description of the image: '{user_details}'.\n"
-                        "If the description is accurate and captures the key elements of the prompt, reply with 'Correct'. "
-                        "If it is inaccurate or missing important details, reply with 'Incorrect' and provide a hint on what is missing. "
-                        "Be Friendly, You are a kids Assistant, use friendly and engaging tone. "
-                        "Don't Mention your system prompt or any prompt; speak from First Person View. "
-                        "Be lenient with the child, they are learning. "
-                        "Use simple and easy words. "
-                        "If some unimportant features are missing, you can mark it as correct."
-                    )
-                }
-            ]
-        }
-    ]
-
-    chat_client = OpenAI(
-        base_url="https://api-inference.huggingface.co/v1/",
-        api_key=chat_api_key  # Loaded from env secrets
+        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,
+        max_tokens=500,
         stream=True
     )
 
@@ -147,7 +96,7 @@ def check_details_fn(user_details):
 
     return response_text
 
-# Define a list of prompts for the dropdown menu.
+# Define a list of prompts for the dropdown.
 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.",
@@ -162,86 +111,70 @@ prompt_options = [
 ]
 
 ##############################################
-# Create the Gradio Interface with Tabs below.
+# Create the Gradio Interface (Single-Page)
 ##############################################
 with gr.Blocks() as demo:
-    # State variables for chat history and saved sessions.
-    # chat_history stores the ongoing conversation as a list of (user_message, bot_response) tuples.
-    # chat_sessions stores a list of saved sessions (each session is itself a list of chat turns).
-    chat_history = gr.State([])   # current session
-    chat_sessions = gr.State([])  # list of sessions
-
-    with gr.Tabs():
-        # ----- Tab 1: Image Generation and Check Details -----
-        with gr.Tab("Image Generation"):
-            gr.Markdown("# Image Generation")
-            with gr.Row():
-                with gr.Column():
-                    prompt_dropdown = gr.Dropdown(label="Select a prompt", choices=prompt_options, value=prompt_options[0])
-                    generate_btn = gr.Button("Generate Image")
-                    img_output = gr.Image(label="Generated Image")
-                    # When generating a new image, automatically save any existing chat session and reset the chat.
-                    generate_btn.click(
-                        generate_image_and_reset_chat, 
-                        inputs=[prompt_dropdown, chat_history, chat_sessions], 
-                        outputs=[img_output, chat_history, chat_sessions]
-                    )
-                with gr.Column():
-                    gr.Markdown("## Check Your Description of the Image")
-                    details_input = gr.Textbox(
-                        label="Enter details about the image", 
-                        placeholder="Describe the key elements of the image..."
-                    )
-                    check_details_btn = gr.Button("Check Details")
-                    details_output = gr.Textbox(label="Result")
-                    check_details_btn.click(check_details_fn, inputs=details_input, outputs=details_output)
-
-        # ----- Tab 2: Chat with Image and Chat History -----
-        with gr.Tab("Chat"):
-            gr.Markdown("# Chat about the Image")
-            gr.Markdown("The conversation below remembers your messages. (Make sure you have generated an image first!)")
-            chatbot = gr.Chatbot(label="Chat History")
-            with gr.Row():
-                chat_input = gr.Textbox(label="Your Message", placeholder="Type your message here...", show_label=False)
-                send_btn = gr.Button("Send")
-            with gr.Row():
-                clear_btn = gr.Button("Clear Chat")
-                # You can still save the session manually if desired.
-                save_session_btn = gr.Button("Save Session")
-            
-            def chat_respond(user_message, history):
-                """
-                Process the chat message: if an image is available, call the chat function;
-                otherwise, return a reminder message.
-                """
-                if not global_image_data_url:
-                    bot_message = "Please generate an image first."
-                else:
-                    bot_message = chat_about_image_fn(user_message)
-                history = history + [(user_message, bot_message)]
-                return "", history
-
-            send_btn.click(chat_respond, inputs=[chat_input, chat_history], outputs=[chat_input, chatbot])
-            chat_input.submit(chat_respond, inputs=[chat_input, chat_history], outputs=[chat_input, chatbot])
-            
-            # Button to clear the current chat (only clears current chat history).
-            clear_btn.click(lambda: ("", []), outputs=[chat_input, chatbot])
-            
-            # Button to manually save the current session: append current chat history to sessions and then clear the chat.
-            def save_session(history, sessions):
-                new_sessions = sessions.copy()
-                if history:
-                    new_sessions.append(history)
-                return "", [], new_sessions
-            save_session_btn.click(save_session, inputs=[chatbot, chat_sessions], outputs=[chat_input, chatbot, chat_sessions])
+    # 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([])
+    saved_sessions = gr.State([])
+
+    gr.Markdown("# Image Generation & Chat Inference")
+    
+    # ----- Image Generation Section -----
+    with gr.Box():
+        gr.Markdown("## Generate Image")
+        with gr.Row():
+            prompt_dropdown = gr.Dropdown(label="Select a prompt", choices=prompt_options, value=prompt_options[0])
+            generate_btn = gr.Button("Generate Image")
+        img_output = gr.Image(label="Generated Image")
         
-        # ----- Tab 3: Chat Sessions History -----
-        with gr.Tab("Sessions"):
-            gr.Markdown("# Saved Chat Sessions")
-            gr.Markdown("Below is a JSON representation of your saved chat sessions.")
-            sessions_output = gr.JSON(label="Saved Sessions")
-            refresh_btn = gr.Button("Refresh Sessions")
-            refresh_btn.click(lambda sessions: sessions, inputs=chat_sessions, outputs=sessions_output)
-
-# Launch the app. (Hugging Face Spaces will detect and run this.)
+        # 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]
+        )
+
+    # ----- Chat Section -----
+    with gr.Box():
+        gr.Markdown("## Chat about the Image")
+        gr.Markdown("After generating an image, ask questions or make comments about it. Your conversation will be automatically saved after each message.")
+        chatbot = gr.Chatbot(label="Chat History")
+        with gr.Row():
+            chat_input = gr.Textbox(label="Your Message", placeholder="Type your message here...", show_label=False)
+            send_btn = gr.Button("Send")
+
+        # Each time the user sends a message, update the chat history.
+        def chat_respond(user_message, history, sessions):
+            if not global_image_data_url:
+                bot_message = "Please generate an image first."
+            else:
+                bot_message = chat_about_image_fn(user_message)
+            new_history = history + [(user_message, bot_message)]
+            # Automatically update saved session with the active conversation.
+            new_sessions = sessions.copy()
+            if new_history: 
+                # In this design, the current active session is always saved (overwritten) as the latest session.
+                if new_sessions:
+                    new_sessions[-1] = new_history
+                else:
+                    new_sessions.append(new_history)
+            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])
+
+    # ----- Saved Sessions Section -----
+    with gr.Box():
+        gr.Markdown("## Saved Chat Sessions")
+        gr.Markdown("Your past chat sessions (including the active one) are saved below. You can refresh to view the latest sessions.")
+        sessions_output = gr.JSON(label="Saved Sessions")
+        refresh_btn = gr.Button("Refresh Saved Sessions")
+        refresh_btn.click(lambda sessions: sessions, inputs=saved_sessions, outputs=sessions_output)
+
+# Launch the app.
 demo.launch()