Update app.py

app.py

@@ -1,8 +1,7 @@
 import torch
-from transformers import AutoProcessor, AutoModelForCausalLM, GenerationConfig, TextIteratorStreamer
+from transformers import AutoProcessor, AutoModelForCausalLM, GenerationConfig
 from PIL import Image
 import gradio as gr
-from threading import Thread
 import spaces
 
 # --- 1. Model and Processor Setup ---
@@ -38,22 +37,32 @@ chat_template = """{% for message in messages -%}
         {%- endif %}"""
 processor.tokenizer.chat_template = chat_template
 
-# --- 2. Gradio Chatbot Logic with Streaming ---
+# --- 2. Gradio Chatbot Logic ---
 @spaces.GPU
-def process_chat_streaming(user_message, chatbot_display, messages_list, image_pil):
+def process_chat(user_message, chatbot_display, messages_list, image_pil):
     """
-    This generator function handles the chat logic with streaming.
-    It yields the updated chatbot display at each step.
+    This function handles the chat logic for a single turn.
+
+    Args:
+        user_message (str): The new message from the user.
+        chatbot_display (list): The current state of the Gradio chatbot display.
+        messages_list (list): The conversation history in the format for the model.
+        image_pil (PIL.Image): The uploaded image.
+
+    Returns:
+        tuple: Updated chatbot_display, updated messages_list, and an empty string for the textbox.
     """
     # Check if an image has been uploaded
     if image_pil is None:
+        # Update the chatbot display with an error message
         chatbot_display.append((user_message, "Please upload an image first to start the conversation."))
-        yield chatbot_display, messages_list
-        return # Stop the generator
+        return chatbot_display, messages_list, "" # Clear the input box
 
-    # Append user's message to the conversation history and display
+    # Append user's message to the conversation history for the model
     messages_list.append({"role": "user", "content": user_message})
-    chatbot_display.append((user_message, "")) # Add an empty spot for the streaming response
+
+    # Append user's message to the chatbot display list
+    chatbot_display.append((user_message, None))
 
     try:
         # Use the processor to apply the chat template
@@ -64,61 +73,52 @@ def process_chat_streaming(user_message, chatbot_display, messages_list, image_p
         )
 
         # Preprocess image and the entire formatted prompt
+        # Patram expects a single image and the full text prompt
         inputs = processor.process(images=[image_pil], text=prompt)
         inputs = {k: v.to(device).unsqueeze(0) for k, v in inputs.items()}
 
-        # Setup the streamer
-        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-        # Define generation configuration
-        generation_config = GenerationConfig(
-            max_new_tokens=512,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.6,
-            stop_strings=["<|endoftext|>", "User:"] # Add stop strings to prevent over-generation
+        # Generate output using model's specific method
+        output = model.generate_from_batch(
+            inputs,
+            GenerationConfig(max_new_tokens=512, do_sample=True, top_p=0.9, temperature=0.6, stop_strings="<|endoftext|>"),
+            tokenizer=processor.tokenizer
         )
 
-        # *** THE FIX IS HERE ***
-        # We must pass 'inputs' as a positional argument for 'batch'
-        # and the rest as keyword arguments.
-        thread = Thread(
-            target=model.generate_from_batch,
-            args=[inputs],  # Pass `inputs` as the first positional argument ('batch')
-            kwargs={       # Pass the rest as keyword arguments
-                "generation_config": generation_config,
-                "tokenizer": processor.tokenizer,
-                "streamer": streamer,
-            }
-        )
-        thread.start()
+        # Extract generated tokens (excluding input tokens) and decode
+        generated_tokens = output[0, inputs['input_ids'].size(1):]
+        response = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
-        # Yield updates to the Gradio UI
-        full_response = ""
-        for new_text in streamer:
-            full_response += new_text
-            chatbot_display[-1] = (user_message, full_response)
-            yield chatbot_display, messages_list
+        # Append assistant's response to the conversation history
+        messages_list.append({"role": "assistant", "content": response})
 
-        # After the loop, the generation is complete.
-        # Add the final full response to the messages list for context.
-        messages_list.append({"role": "assistant", "content": full_response})
-        yield chatbot_display, messages_list # Yield the final state
+        # Update the chatbot display with the assistant's response
+        chatbot_display[-1] = (user_message, response)
 
     except Exception as e:
-        print(f"Error during streaming inference: {e}")
-        error_message = f"Sorry, an error occurred: {e}"
+        print(f"Error during inference: {e}")
+        error_message = f"Sorry, an error occurred during processing: {e}"
+        # Update the last message in the chatbot display with the error
         chatbot_display[-1] = (user_message, error_message)
-        yield chatbot_display, messages_list
+
+    # Return the updated state and clear the input textbox
+    return chatbot_display, messages_list, ""
+
+
+def clear_chat(chatbot_display, messages_list, image_input):
+    """Resets the chat, history, and image."""
+    return [], [], None, "Type your question here..."
+
 
 # --- 3. Gradio Interface Definition ---
 
 with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutral")) as demo:
-    gr.Markdown("# 🤖 Patram-7B-Instruct Streaming Chatbot")
-    gr.Markdown("Upload an image and ask questions about it. The response will stream in real-time.")
+    gr.Markdown("# 🤖 Patram-7B-Instruct Chatbot")
+    gr.Markdown("Upload an image and ask questions about it. The chatbot will remember the conversation context.")
 
-    # State variables to hold conversation history
+    # State variables to hold conversation history and image
     messages_list = gr.State([])
+    # We don't need a state for chatbot_display as it's passed as an input/output directly
+    # The image is also passed directly from the gr.Image component
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -129,8 +129,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutra
             chatbot_display = gr.Chatbot(
                 label="Conversation",
                 bubble_full_width=False,
-                height=500,
-                avatar_images=(None, "https://cdn-avatars.huggingface.co/v1/production/uploads/67b462a1f4f414c2b3e2bc2f/EnVeNWEIeZ6yF6ueZ7E3Y.jpeg")
+                height=500
             )
             with gr.Row():
                 user_textbox = gr.Textbox(
@@ -139,22 +138,21 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutra
                     scale=4,
                     container=False
                 )
+                submit_btn = gr.Button("Send", variant="primary", scale=1, min_width=0)
+
 
     # --- Event Listeners ---
 
-    # Define the action for submitting a message (via enter key)
+    # Define the action for submitting a message (via button or enter key)
     submit_action = user_textbox.submit(
-        fn=process_chat_streaming,
+        fn=process_chat,
         inputs=[user_textbox, chatbot_display, messages_list, image_input],
-        outputs=[chatbot_display, messages_list],
+        outputs=[chatbot_display, messages_list, user_textbox]
     )
-
-    # Chain the action to also clear the textbox after submission
-    submit_action.then(
-        fn=lambda: gr.update(value=""),
-        inputs=None,
-        outputs=[user_textbox],
-        queue=False
+    submit_btn.click(
+        fn=process_chat,
+        inputs=[user_textbox, chatbot_display, messages_list, image_input],
+        outputs=[chatbot_display, messages_list, user_textbox]
     )
 
     # Define the action for the clear button