import gradio as gr import os import zipfile from io import BytesIO import time import tempfile from pathlib import Path import shutil from main import process_images from prompt import optimize_prompt # Maximum number of images MAX_IMAGES = 30 # ------- File Operations ------- def create_download_file(image_paths, captions): """Create a zip file with images and their captions""" zip_io = BytesIO() with zipfile.ZipFile(zip_io, 'w') as zip_file: for i, (image_path, caption) in enumerate(zip(image_paths, captions)): # Get original filename without extension base_name = os.path.splitext(os.path.basename(image_path))[0] img_name = f"{base_name}.png" caption_name = f"{base_name}.txt" # Add image to zip with open(image_path, 'rb') as img_file: zip_file.writestr(img_name, img_file.read()) # Add caption to zip zip_file.writestr(caption_name, caption) return zip_io.getvalue() def process_uploaded_images(image_paths, batch_by_category=False): """Process uploaded images using main.py's functions""" # Create temporary directories for input and output with tempfile.TemporaryDirectory() as temp_input_dir, tempfile.TemporaryDirectory() as temp_output_dir: # Copy all images to the temporary input directory temp_input_path = Path(temp_input_dir) temp_output_path = Path(temp_output_dir) # Map of original paths to filenames in temp dir path_mapping = {} for i, path in enumerate(image_paths): # Keep original filename to preserve categorization filename = os.path.basename(path) temp_path = temp_input_path / filename # Copy file to temp directory shutil.copy2(path, temp_path) path_mapping[str(temp_path)] = str(path) # Process the images using main.py's function process_images(temp_input_dir, temp_output_dir, batch_images=batch_by_category) # Collect the captions from the output directory captions = [] for orig_path in image_paths: # Get the base filename without extension base_name = os.path.splitext(os.path.basename(orig_path))[0] caption_filename = f"{base_name}.txt" caption_path = temp_output_path / caption_filename # If caption file exists, read it; otherwise use empty string if os.path.exists(caption_path): with open(caption_path, 'r', encoding='utf-8') as f: caption = f.read() captions.append(caption) else: captions.append("") return captions # ------- UI Interaction Functions ------- def load_captioning(files): """Process uploaded images and show them in the UI""" if not files: return [], gr.update(visible=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(visible=False), gr.update(value="Upload images to begin"), *[gr.update(visible=False) for _ in range(MAX_IMAGES)] # Filter to only keep image files image_paths = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'))] if not image_paths or len(image_paths) < 1: gr.Warning(f"Please upload at least one image") return [], gr.update(visible=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(visible=False), gr.update(value="No valid images found"), *[gr.update(visible=False) for _ in range(MAX_IMAGES)] if len(image_paths) > MAX_IMAGES: gr.Warning(f"Only the first {MAX_IMAGES} images will be processed") image_paths = image_paths[:MAX_IMAGES] # Update row visibility row_updates = [] for i in range(MAX_IMAGES): if i < len(image_paths): row_updates.append(gr.update(visible=True)) else: row_updates.append(gr.update(visible=False)) return ( image_paths, # stored_image_paths gr.update(visible=True), # captioning_area gr.update(interactive=True), # caption_btn gr.update(interactive=False), # download_btn - initially disabled until captioning is done gr.update(visible=False), # download_output gr.update(value=f"{len(image_paths)} images ready for captioning"), # status_text *row_updates # image_rows ) def update_images(image_paths): """Update the image components with the uploaded images""" print(f"Updating images with paths: {image_paths}") updates = [] for i in range(MAX_IMAGES): if i < len(image_paths): updates.append(gr.update(value=image_paths[i])) else: updates.append(gr.update(value=None)) return updates def update_caption_labels(image_paths): """Update caption labels to include the image filename""" updates = [] for i in range(MAX_IMAGES): if i < len(image_paths): filename = os.path.basename(image_paths[i]) updates.append(gr.update(label=filename)) else: updates.append(gr.update(label="")) return updates def run_captioning(image_paths, batch_category): """Generate captions for the images using the main.py functions""" if not image_paths: return [gr.update(value="") for _ in range(MAX_IMAGES)] + [gr.update(value="No images to process")] try: print(f"Starting captioning for {len(image_paths)} images, batch_by_category={batch_category}") captions = process_uploaded_images(image_paths, batch_category) # Count valid captions valid_captions = sum(1 for c in captions if c and c.strip()) print(f"Generated {valid_captions} valid captions out of {len(captions)} images") if valid_captions < len(captions): gr.Warning(f"{len(captions) - valid_captions} images could not be captioned properly") status = gr.update(value=f"✅ Captioning complete - {valid_captions}/{len(captions)} successful") else: gr.Info("Captioning complete!") status = gr.update(value="✅ Captioning complete") print("Sample captions:", captions[:2] if len(captions) >= 2 else captions) except Exception as e: print(f"Error in captioning: {str(e)}") gr.Error(f"Captioning failed: {str(e)}") captions = [""] * len(image_paths) # Use empty strings status = gr.update(value=f"❌ Error: {str(e)}") # Update caption textboxes caption_updates = [] for i in range(MAX_IMAGES): if i < len(captions) and captions[i]: # Only set value if we have a valid caption caption_updates.append(gr.update(value=captions[i])) else: caption_updates.append(gr.update(value="")) print(f"Returning {len(caption_updates)} caption updates") return caption_updates + [status] def update_batch_setting(value): """Update the batch by category setting""" return value def create_zip_from_ui(image_paths, *captions_list): """Create a zip file from the current images and captions in the UI""" # Filter out empty captions for non-existent images valid_captions = [cap for i, cap in enumerate(captions_list) if i < len(image_paths) and cap] valid_image_paths = image_paths[:len(valid_captions)] if not valid_image_paths: gr.Warning("No images to download") return None # Create zip file zip_data = create_download_file(valid_image_paths, valid_captions) timestamp = time.strftime("%Y%m%d_%H%M%S") # Create a temporary file to store the zip temp_dir = tempfile.gettempdir() zip_filename = f"image_captions_{timestamp}.zip" zip_path = os.path.join(temp_dir, zip_filename) # Write the zip data to the temporary file with open(zip_path, "wb") as f: f.write(zip_data) # Return the path to the temporary file return zip_path def process_upload(files, image_rows, image_components, caption_components): """Process uploaded files and update UI components""" # First get paths and visibility updates image_paths, captioning_update, caption_btn_update, download_btn_update, download_output_update, status_update, *row_updates = load_captioning(files) # Then get image updates image_updates = update_images(image_paths) # Update caption labels with filenames caption_label_updates = update_caption_labels(image_paths) # Return all updates together return [image_paths, captioning_update, caption_btn_update, download_btn_update, download_output_update, status_update] + row_updates + image_updates + caption_label_updates def on_captioning_start(): """Update UI when captioning starts""" return gr.update(value="⏳ Processing captions... please wait"), gr.update(interactive=False) def on_captioning_complete(): """Update UI when captioning completes""" return gr.update(value="✅ Captioning complete"), gr.update(interactive=True), gr.update(interactive=True) # ------- UI Style Definitions ------- def get_css_styles(): """Return CSS styles for the UI""" return """ """ # ------- UI Component Creation ------- def create_upload_area(): """Create the upload area components""" # Left column for images/upload with gr.Column(scale=1, elem_id="left-column") as upload_column: # Upload area gr.Markdown("### Upload your images", elem_id="upload-heading") gr.Markdown("Only .png, .jpg, .jpeg, and .webp files are supported", elem_id="file-types-info", elem_classes="file-types-info") image_upload = gr.File( file_count="multiple", label="Drop your files here", file_types=["image"], type="filepath", height=220, elem_classes="file-upload-container", ) return upload_column, image_upload def create_config_area(): """Create the configuration area components""" # Right column for configuration and captions with gr.Column(scale=1.5, elem_id="right-column") as config_column: # Configuration area gr.Markdown("### Configuration") batch_category_checkbox = gr.Checkbox( label="Batch process by category", value=False, info="Caption similar images together" ) gr.Markdown(""" **Note about categorization:** - Images are grouped by the part of the filename before the last underscore - For example: "character_pose_1.png" and "character_pose_2.png" share the category "character_pose" - When using "Batch process by category", similar images are captioned together for more consistent results """, elem_classes=["category-info"]) caption_btn = gr.Button("Caption Images", variant="primary", interactive=False) download_btn = gr.Button("Download Images + Captions", variant="secondary", interactive=False) download_output = gr.File(label="Download Zip", visible=False) status_text = gr.Markdown("Upload images to begin", visible=True) return config_column, batch_category_checkbox, caption_btn, download_btn, download_output, status_text def create_captioning_area(): """Create the captioning area components""" with gr.Column(visible=False) as captioning_area: # Replace the single heading with a row containing two headings with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Your Images", elem_id="images-heading") with gr.Column(scale=1.5): gr.Markdown("### Your Captions", elem_id="captions-heading") # Create individual image and caption rows image_rows = [] image_components = [] caption_components = [] for i in range(MAX_IMAGES): with gr.Row(visible=False, elem_classes=["image-caption-row"]) as img_row: image_rows.append(img_row) # Left column for image with gr.Column(scale=1): img = gr.Image( label=f"Image {i+1}", type="filepath", show_label=False, height=200, width=200, elem_classes=["image-thumbnail"] ) image_components.append(img) # Right column for caption with gr.Column(scale=1.5): caption = gr.Textbox( label=f"Caption {i+1}", lines=3, elem_classes=["caption-area"] ) caption_components.append(caption) return captioning_area, image_rows, image_components, caption_components def setup_event_handlers( image_upload, stored_image_paths, captioning_area, caption_btn, download_btn, download_output, status_text, image_rows, image_components, caption_components, batch_category_checkbox, batch_by_category, shared_captions=None ): """Set up all event handlers for the UI""" # Combined outputs for the upload function upload_outputs = [ stored_image_paths, captioning_area, caption_btn, download_btn, download_output, status_text, *image_rows ] combined_outputs = upload_outputs + image_components + caption_components # Set up upload handler image_upload.change( lambda files: process_upload(files, image_rows, image_components, caption_components), inputs=[image_upload], outputs=combined_outputs ) # Set up batch category checkbox batch_category_checkbox.change( update_batch_setting, inputs=[batch_category_checkbox], outputs=[batch_by_category] ) # Set up captioning button chain caption_chain = caption_btn.click( on_captioning_start, inputs=None, outputs=[status_text, caption_btn] ).success( run_captioning, inputs=[stored_image_paths, batch_by_category], outputs=caption_components + [status_text] ).success( on_captioning_complete, inputs=None, outputs=[status_text, caption_btn, download_btn] ) # If shared_captions is provided, add an additional handler to update it if shared_captions is not None: def extract_valid_captions(*caption_values): return [c for c in caption_values if c and c.strip()] caption_chain.success( extract_valid_captions, inputs=caption_components, outputs=[shared_captions] ) # Set up download button download_btn.click( create_zip_from_ui, inputs=[stored_image_paths] + caption_components, outputs=[download_output] ).then( lambda: gr.update(visible=True, elem_classes=["download-section"]), inputs=None, outputs=[download_output] ).then( lambda: gr.Info("Click the Download button that appeared to save your zip file"), inputs=None, outputs=None ) # ------- Prompt Optimization UI ------- def create_prompt_optimization_ui(): """Create UI components for prompt optimization tab""" with gr.Column(scale=1) as left_column: # Left side for caption input gr.Markdown("### Upload Captions") gr.Markdown("Upload caption files (.txt) or enter captions manually", elem_classes="file-types-info") captions_upload = gr.File( file_count="multiple", label="Upload caption files", file_types=[".txt"], type="filepath", elem_classes="file-upload-container", height=220 ) manual_captions = gr.Textbox( label="Or enter captions manually", lines=5, placeholder="Enter captions here, one per line", elem_classes="prompt-box" ) # Add button to use captions from image captioning tab use_generated_captions = gr.Button("Use Captions from Manual Entry", variant="secondary") with gr.Column(scale=1) as right_column: # Right side for prompt input and output gr.Markdown("### Optimize Prompt") gr.Markdown("\n- Craft prompts that match the style of your training captions\n- Enter a simple prompt and receive an optimized version\n", elem_classes=["category-info"]) user_prompt = gr.Textbox( label="Enter your prompt", lines=3, placeholder="Enter the prompt you want to optimize", elem_classes="prompt-box" ) optimize_btn = gr.Button("Optimize Prompt", variant="primary", elem_classes="optimize-btn") optimized_prompt = gr.Textbox( label="Optimized Prompt", lines=5, placeholder="Optimized prompt will appear here", elem_classes="prompt-box" ) optimization_status = gr.Markdown("Enter a prompt and upload captions to begin", elem_classes="optimization-status") # Return components but NOT info_md (will create it separately in build_ui) return ( left_column, right_column, captions_upload, manual_captions, use_generated_captions, user_prompt, optimize_btn, optimized_prompt, optimization_status ) def run_optimization(prompt, caption_files, manual_caption_text): """Handle the prompt optimization logic""" if not prompt or prompt.strip() == "": return "", "Please enter a prompt to optimize" # Handle different input sources for captions caption_list = [] if manual_caption_text and manual_caption_text.strip(): # Use manually entered captions caption_list = [line.strip() for line in manual_caption_text.split("\n") if line.strip()] elif caption_files and len(caption_files) > 0: # Read captions from uploaded files for file_path in caption_files: if os.path.exists(file_path) and file_path.lower().endswith('.txt'): with open(file_path, 'r', encoding='utf-8') as f: content = f.read().strip() if content: caption_list.append(content) if not caption_list: return "", "Please upload caption files or enter captions manually" try: # Call the optimize_prompt function from prompt.py result = optimize_prompt(prompt, captions_list=caption_list) return result, "✅ Prompt optimization complete" except Exception as e: return "", f"❌ Error optimizing prompt: {str(e)}" def setup_prompt_optimization_handlers( captions_upload, manual_captions, use_generated_captions, user_prompt, optimize_btn, optimized_prompt, optimization_status, shared_captions ): """Set up event handlers for prompt optimization tab""" # Function to update manual captions with shared ones def fill_with_shared_captions(captions_list): if not captions_list or len(captions_list) == 0: return "No captions available. Generate captions in the Image Captioning tab first." return "\n".join(captions_list) # Connect button to fill manual captions area use_generated_captions.click( fill_with_shared_captions, inputs=[shared_captions], outputs=[manual_captions] ) # Connect the optimize button to the optimization function optimize_btn.click( run_optimization, inputs=[user_prompt, captions_upload, manual_captions], outputs=[optimized_prompt, optimization_status] ) # ------- Main Application ------- def build_ui(): """Build and return the Gradio interface""" with gr.Blocks() as demo: gr.Markdown("# Image Auto-captioner for LoRA Training") gr.Markdown("""Check out the [code](https://github.com/RishiDesai/LoRACaptioner) and see my [blog post](https://rishidesai.github.io/posts/character-lora/) for more information.""") # Store generated captions for sharing between tabs shared_captions = gr.State([]) # Create tabs for different functionality with gr.Tabs() as tabs: with gr.TabItem("Image Captioning") as captioning_tab: # Store uploaded images stored_image_paths = gr.State([]) batch_by_category = gr.State(False) # State to track if batch by category is enabled # Create a two-column layout for the entire interface with gr.Row(): # Create upload area in left column _, image_upload = create_upload_area() # Create config area in right column _, batch_category_checkbox, caption_btn, download_btn, download_output, status_text = create_config_area() # Create captioning area (initially hidden) captioning_area, image_rows, image_components, caption_components = create_captioning_area() # Set up event handlers with shared captions setup_event_handlers( image_upload, stored_image_paths, captioning_area, caption_btn, download_btn, download_output, status_text, image_rows, image_components, caption_components, batch_category_checkbox, batch_by_category, shared_captions ) with gr.TabItem("Prompt Optimization") as prompt_tab: with gr.Row(): # Create prompt optimization UI components ( left_column, right_column, captions_upload, manual_captions, use_generated_captions, user_prompt, optimize_btn, optimized_prompt, optimization_status ) = create_prompt_optimization_ui() # Set up prompt optimization event handlers setup_prompt_optimization_handlers( captions_upload, manual_captions, use_generated_captions, user_prompt, optimize_btn, optimized_prompt, optimization_status, shared_captions ) # Add CSS styling gr.HTML(get_css_styles()) return demo # Launch the app when run directly if __name__ == "__main__": demo = build_ui() demo.launch(share=True)