Update app.py

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 import warnings
 import gradio as gr
 import os
+from gradio_client import Client
 
 # Disable warnings for cleaner output
 transformers.logging.set_verbosity_error()
@@ -32,6 +33,10 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 print("Model loaded successfully!")
 
+# Initialize the client for the test bot
+chatter = "K00B404/transcript_image_generator"
+chatbot_client = Client(chatter)
+
 def analyze_character(image_path, analysis_type):
     """
     Analyze a character image for dramaturgical insights
@@ -46,6 +51,9 @@ def analyze_character(image_path, analysis_type):
     # Load and process image
     try:
         image = Image.open(image_path).convert('RGB')
+        # Resize image to 512x512
+        image = image.resize((256, 256), Image.Resampling.LANCZOS)
+        
         image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
     except Exception as e:
         return f"Error processing image: {str(e)}"
@@ -118,30 +126,132 @@ def analyze_character(image_path, analysis_type):
         except Exception as e2:
             return f"Error generating analysis: {str(e)}\nFallback also failed: {str(e2)}\n\nPlease try a different image or check model compatibility."
 
+def chat_with_persona(message, history, system_message, max_tokens, temperature, top_p):
+    """Function to interact with the chatbot API using the generated persona"""
+    try:
+        # Call the API with the current message and system prompt (persona)
+        response = chatbot_client.predict(
+            message=message,
+            system_message=system_message,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            api_name="/chat"
+        )
+        return response
+    except Exception as e:
+        return f"Error communicating with the chatbot API: {str(e)}"
+
 # Create Gradio interface
 def create_ui():
     with gr.Blocks(title="Dramaturg Character Analyzer") as demo:
-        gr.Markdown("# Dramaturg Character Analyzer")
-        gr.Markdown("Upload a character image to receive a dramaturgical analysis")
+        # Store the current analysis result for sharing between tabs
+        analysis_result = gr.State("")
         
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(type="filepath", label="Upload Character Image")
-                analysis_type = gr.Radio(
-                    ["full_analysis", "archetype", "historical_context", "basic_description"],
-                    label="Analysis Type",
-                    value="full_analysis"
+        with gr.Tabs() as tabs:
+            # First tab: Character analysis
+            with gr.TabItem("Character Analysis"):
+                gr.Markdown("# Dramaturg Character Analyzer")
+                gr.Markdown("Upload a character image to receive a dramaturgical analysis")
+                
+                with gr.Row():
+                    with gr.Column():
+                        input_image = gr.Image(type="filepath", label="Upload Character Image")
+                        analysis_type = gr.Radio(
+                            ["full_analysis", "archetype", "historical_context", "basic_description"],
+                            label="Analysis Type",
+                            value="full_analysis"
+                        )
+                        analyze_btn = gr.Button("Analyze Character")
+                    
+                    with gr.Column():
+                        output_text = gr.Textbox(label="Character Analysis", lines=20)
+                        copy_to_test_btn = gr.Button("Copy to Test Bot", interactive=False)
+                
+                def update_analysis_result(result):
+                    # Enable the copy button once we have a result
+                    return result, True
+                
+                analyze_btn.click(
+                    fn=analyze_character,
+                    inputs=[input_image, analysis_type],
+                    outputs=[output_text, copy_to_test_btn]
+                )
+                
+                def copy_to_test(result):
+                    # Update the state and switch to the test tab
+                    return result, 1
+                
+                copy_to_test_btn.click(
+                    fn=copy_to_test,
+                    inputs=[output_text],
+                    outputs=[analysis_result, tabs]
                 )
-                analyze_btn = gr.Button("Analyze Character")
             
-            with gr.Column():
-                output_text = gr.Textbox(label="Character Analysis", lines=20)
-        
-        analyze_btn.click(
-            fn=analyze_character,
-            inputs=[input_image, analysis_type],
-            outputs=output_text
-        )
+            # Second tab: Test bot integration
+            with gr.TabItem("Test Bot"):
+                gr.Markdown("# Test Your Character Persona")
+                gr.Markdown("The character analysis will be used as the system prompt for the test bot.")
+                
+                with gr.Row():
+                    with gr.Column():
+                        system_prompt = gr.Textbox(label="System Prompt (Character Persona)", lines=10)
+                        
+                        with gr.Row():
+                            max_tokens = gr.Slider(minimum=100, maximum=4000, value=1000, step=100, label="Max Tokens")
+                            temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+                            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top P")
+                        
+                        user_input = gr.Textbox(label="Your message", placeholder="Ask something about the character...")
+                        send_btn = gr.Button("Send Message")
+                    
+                    with gr.Column():
+                        chatbot = gr.Chatbot(label="Conversation")
+                
+                def update_system_prompt(result):
+                    return result
+                
+                # Update the system prompt when switching to this tab with an analysis result
+                demo.load(
+                    fn=update_system_prompt,
+                    inputs=[analysis_result],
+                    outputs=[system_prompt]
+                )
+                
+                # Chat history for the test bot
+                chat_history = []
+                
+                def respond(message, history, system_message, max_tokens_val, temperature_val, top_p_val):
+                    # Add the user message to history
+                    history.append((message, ""))
+                    
+                    # Get response from the test bot
+                    response = chat_with_persona(
+                        message=message,
+                        history=history,
+                        system_message=system_message,
+                        max_tokens=max_tokens_val,
+                        temperature=temperature_val,
+                        top_p=top_p_val
+                    )
+                    
+                    # Update the last history item with the response
+                    history[-1] = (message, response)
+                    
+                    return "", history
+                
+                send_btn.click(
+                    fn=respond,
+                    inputs=[user_input, chatbot, system_prompt, max_tokens, temperature, top_p],
+                    outputs=[user_input, chatbot]
+                )
+                
+                # Also trigger on pressing Enter in the input box
+                user_input.submit(
+                    fn=respond,
+                    inputs=[user_input, chatbot, system_prompt, max_tokens, temperature, top_p],
+                    outputs=[user_input, chatbot]
+                )
         
     return demo
 
@@ -149,4 +259,4 @@ def create_ui():
 if __name__ == "__main__":
     demo = create_ui()
     demo.launch(share=True)
-    print("Dramaturg Character Analyzer is now running!")
+    print("Dramaturg Character Analyzer is now running with Test Bot integration!")