Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 from PIL import Image
 import os 
 import io
+import glob
 from transformers import BlipProcessor, BlipForConditionalGeneration
 import time
 from gradio_client import Client
@@ -30,8 +31,25 @@ model = BlipForConditionalGeneration.from_pretrained(blipper)
 chatbot_client = Client(chatter)
 
 PROFILES_FILENAME = "character_profiles.json"
-REPO_ID = "K00B404/logging_space"
+REPO_ID = "K00B404/Persona_from_Image"
 REPO_TYPE = "space"
+CHARACTERS_FOLDER = "characters"  # Root folder for character images
+
+# Ensure characters folder exists
+if not os.path.exists(CHARACTERS_FOLDER):
+    os.makedirs(CHARACTERS_FOLDER)
+
+def get_character_images():
+    """Get all image files from the characters folder"""
+    image_files = []
+    # Get all PNG, JPG, JPEG, and WebP files
+    for ext in ['png', 'jpg', 'jpeg', 'webp']:
+        image_files.extend(glob.glob(f"{CHARACTERS_FOLDER}/*.{ext}"))
+        image_files.extend(glob.glob(f"{CHARACTERS_FOLDER}/*.{ext.upper()}"))
+    
+    # Sort alphabetically
+    image_files.sort()
+    return image_files
 
 def load_profiles():
     """Load profiles from HuggingFace space"""
@@ -155,7 +173,7 @@ def generate_persona(img, min_len, max_len, persona_detail_level):
     # dramaturg to mae a solid role for a actor from pragmatic description
     system_prompt="You are a Expert Dramaturg and your task is to use the input persona information and write a 'Role' description as compact instuctions for the actor"
     persona = helper_llm(persona, system_prompt=system_prompt)
-    return caption, persona, total_time
+    return caption, persona, total_time, img
 
 
 def chat_with_persona(message, history, system_message, max_tokens, temperature, top_p):
@@ -211,11 +229,28 @@ def generate_flux_image(final_prompt, is_negative, steps, cfg_scale, seed, stren
         print(f"Error when trying to open the image: {e}")
         return None, None, None
 
+def process_selected_image(image_path):
+    """Process a selected image from the gallery to generate a persona"""
+    if not image_path:
+        return "", "", "No image selected", None
+    
+    try:
+        # Use default values for generation
+        min_len = 50
+        max_len = 200
+        persona_detail_level = "Comprehensive"
+        
+        caption, persona, time_output, _ = generate_persona(image_path, min_len, max_len, persona_detail_level)
+        return caption, persona, time_output, image_path
+    except Exception as e:
+        return "", "", f"Error processing image: {str(e)}", None
+
 # Create Gradio interface with tabs
 with gr.Blocks(title="Image Character Persona Generator") as iface:
-    # Store the generated persona in a state variable to share between tabs
+    # Store state variables
     persona_state = gr.State("")
     profiles_state = gr.State(load_profiles())
+    selected_image_state = gr.State(None)
     
     with gr.Tabs():
         # First tab: Persona Generator
@@ -321,23 +356,59 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
             
             with gr.Row():
                 with gr.Column(scale=1):
+                    # Character Gallery
+                    gr.Markdown("### Character Images")
+                    
+                    # Auto-load images from characters folder
+                    character_gallery = gr.Gallery(
+                        label="Available Character Images", 
+                        elem_id="character_gallery",
+                        columns=3,
+                        object_fit="contain",
+                        height="auto"
+                    )
+                    
+                    refresh_gallery_btn = gr.Button("Refresh Character Gallery")
+                    
                     # Profile Management Section
+                    gr.Markdown("### Profile Management")
                     profile_name = gr.Textbox(label="Profile Name", placeholder="Enter a name for this character profile")
                     add_profile_btn = gr.Button("Save Current Profile")
                     status_msg = gr.Textbox(label="Status", interactive=False)
                     
                     # Profile Selection Dropdown
-                    profile_dropdown = gr.Dropdown(label="Select Profile", interactive=True)
+                    profile_dropdown = gr.Dropdown(label="Select Saved Profile", interactive=True)
                     load_profile_btn = gr.Button("Load Selected Profile")
                     update_profile_btn = gr.Button("Update Selected Profile")
                     delete_profile_btn = gr.Button("Delete Selected Profile")
                     refresh_profiles_btn = gr.Button("Refresh Profiles List")
                 
                 with gr.Column(scale=2):
+                    # Profile Preview
+                    preview_image = gr.Image(label="Character Image Preview", type="filepath")
+                    
                     # Profile Editor
                     profile_caption = gr.Textbox(label="Character Description", lines=3)
                     profile_persona = gr.Textbox(label="Character Persona", lines=15)
                     image_path_display = gr.Textbox(label="Image Path (Reference Only)", interactive=False)
+                    
+                    # Process Selected Image Button
+                    process_image_btn = gr.Button("Process Selected Image")
+                    
+                    # Copy to Chat Button
+                    copy_to_chat_btn = gr.Button("Use This Persona in Chat")
+            
+            # Function to refresh the character gallery
+            def refresh_character_gallery():
+                image_files = get_character_images()
+                return gr.Gallery(value=image_files)
+            
+            # Function to handle gallery selection
+            def gallery_select(evt: gr.SelectData, gallery_images):
+                if gallery_images and evt.index < len(gallery_images):
+                    selected_image = gallery_images[evt.index]
+                    return selected_image, selected_image
+                return None, None
             
             # Function to refresh the profiles dropdown
             def refresh_profiles_list(profiles_data):
@@ -345,7 +416,7 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                 return gr.Dropdown(choices=[p["name"] for p in profiles], value=None if not profiles else profiles[0]["name"])
             
             # Function to save current profile
-            def save_current_profile(name, caption, persona, input_image_path, profiles_data):
+            def save_current_profile(name, caption, persona, image_path, profiles_data):
                 if not name or not caption or not persona:
                     return profiles_data, "Error: Please provide a name, caption, and persona", gr.Dropdown(choices=[p["name"] for p in profiles_data["profiles"]])
                 
@@ -354,7 +425,7 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                     "name": name,
                     "caption": caption,
                     "persona": persona,
-                    "image_path": input_image_path if input_image_path else "",
+                    "image_path": image_path if image_path else "",
                     "created_at": datetime.now().isoformat()
                 }
                 
@@ -378,12 +449,13 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
             def load_selected_profile(profile_name, profiles_data):
                 for profile in profiles_data["profiles"]:
                     if profile["name"] == profile_name:
-                        return profile["caption"], profile["persona"], profile["image_path"], "Profile loaded successfully", profile["persona"]
+                        image_path = profile["image_path"] if os.path.exists(profile["image_path"]) else None
+                        return profile["caption"], profile["persona"], profile["image_path"], image_path, "Profile loaded successfully", profile["persona"]
                 
-                return "", "", "", "Error: Profile not found", ""
+                return "", "", "", None, "Error: Profile not found", ""
             
             # Function to update selected profile
-            def update_selected_profile(profile_name, caption, persona, profiles_data):
+            def update_selected_profile(profile_name, caption, persona, image_path, profiles_data):
                 if not profile_name:
                     return profiles_data, "Error: No profile selected"
                 
@@ -392,6 +464,8 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                         # Update profile data
                         profiles_data["profiles"][i]["caption"] = caption
                         profiles_data["profiles"][i]["persona"] = persona
+                        if image_path and image_path != profile["image_path"]:
+                            profiles_data["profiles"][i]["image_path"] = image_path
                         profiles_data["profiles"][i]["updated_at"] = datetime.now().isoformat()
                         
                         # Save updated profiles
@@ -415,22 +489,43 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                 profile_names = [p["name"] for p in profiles_data["profiles"]]
                 return profiles_data, f"Deleted profile: {profile_name}", gr.Dropdown(choices=profile_names, value=None if not profile_names else profile_names[0])
             
+            # Function to use the current persona in chat
+            def use_persona_in_chat(persona):
+                return persona
+            
             # Connect the UI elements with their functions
+            refresh_gallery_btn.click(
+                fn=refresh_character_gallery,
+                outputs=[character_gallery]
+            )
+            
+            character_gallery.select(
+                fn=gallery_select,
+                inputs=[character_gallery],
+                outputs=[selected_image_state, preview_image]
+            )
+            
+            process_image_btn.click(
+                fn=process_selected_image,
+                inputs=[selected_image_state],
+                outputs=[profile_caption, profile_persona, status_msg, image_path_display]
+            )
+            
             add_profile_btn.click(
                 fn=save_current_profile,
-                inputs=[profile_name, profile_caption, profile_persona, input_image, profiles_state],
+                inputs=[profile_name, profile_caption, profile_persona, selected_image_state, profiles_state],
                 outputs=[profiles_state, status_msg, profile_dropdown]
             )
             
             load_profile_btn.click(
                 fn=load_selected_profile,
                 inputs=[profile_dropdown, profiles_state],
-                outputs=[profile_caption, profile_persona, image_path_display, status_msg, system_prompt]
+                outputs=[profile_caption, profile_persona, image_path_display, preview_image, status_msg, system_prompt]
             )
             
             update_profile_btn.click(
                 fn=update_selected_profile,
-                inputs=[profile_dropdown, profile_caption, profile_persona, profiles_state],
+                inputs=[profile_dropdown, profile_caption, profile_persona, selected_image_state, profiles_state],
                 outputs=[profiles_state, status_msg]
             )
             
@@ -446,35 +541,41 @@ with gr.Blocks(title="Image Character Persona Generator") as iface:
                 outputs=[profile_dropdown]
             )
             
-            # Initialize the profiles dropdown
+            copy_to_chat_btn.click(
+                fn=use_persona_in_chat,
+                inputs=[profile_persona],
+                outputs=[system_prompt]
+            )
+            
+            # Initialize the profiles dropdown and character gallery
             iface.load(
-                fn=refresh_profiles_list,
+                fn=lambda: (refresh_profiles_list, refresh_character_gallery),
                 inputs=[profiles_state],
-                outputs=[profile_dropdown]
+                outputs=[profile_dropdown, character_gallery]
             )
     
     # Function to update system prompt in Test tab when persona is generated
-    def update_persona_state(caption, persona, time_output):
+    def update_persona_state(caption, persona, time_output, img_path):
         return persona, persona
     
     # Connect the persona generator to update the system prompt
     submit_btn.click(fn=generate_persona, 
                     inputs=[input_image, min_length, max_length, detail_level],
-                    outputs=[caption_output, persona_output, time_output])
+                    outputs=[caption_output, persona_output, time_output, selected_image_state])
     
     # Update the system prompt in Test tab when persona is generated
     submit_btn.click(fn=update_persona_state,
-                    inputs=[caption_output, persona_output, time_output],
+                    inputs=[caption_output, persona_output, time_output, input_image],
                     outputs=[persona_state, system_prompt])
     
     # Function to update profile fields when a new persona is generated
-    def update_profile_fields(caption, persona):
-        return caption, persona
+    def update_profile_fields(caption, persona, img_path):
+        return caption, persona, img_path, img_path
     
     # Update the profile fields in Persona Manager tab when persona is generated
     submit_btn.click(fn=update_profile_fields,
-                    inputs=[caption_output, persona_output],
-                    outputs=[profile_caption, profile_persona])
+                    inputs=[caption_output, persona_output, input_image],
+                    outputs=[profile_caption, profile_persona, selected_image_state, preview_image])
 
 # Launch the interface
 iface.launch()