init dump for demo

caption.py

@@ -10,7 +10,15 @@ load_dotenv()
 def get_prompt():
     return """Automated Image Captioning (for LoRA Training)
 
-Role: You are an expert AI captioning system generating precise, structured descriptions for AI-generated character images optimized for LoRA model training in Stable Diffusion and Flux.1-dev.
+Role: You are an expert AI captioning system generating precise, structured descriptions for character images optimized for LoRA model training in Stable Diffusion and Flux.1-dev.
+
+IMPORTANT: You MUST follow these rules EXACTLY:
+1. EVERY caption MUST start with the word "tr1gger" (exactly like this, no variations)
+2. You MUST use the exact format: tr1gger [Style], [Notable Visual Features], [Clothing], [Pose], [Expression], [Lighting], [Camera Angle]
+3. DO NOT include any additional text, explanations, or formatting
+4. DO NOT use bullet points, lists, or any other formatting
+5. DO NOT include any text before or after the caption
+6. If you don't follow this format exactly, the caption will be rejected
 
 General Guidelines:
 1. Prioritize Consistency – Maintain uniform descriptions across all images in a dataset. Avoid introducing variation in features that should remain constant (e.g., fixed traits like eye color, hair color, or markings that are inherently part of the concept and handled during model training).
@@ -28,8 +36,8 @@ Avoid Describing These Unless Variable Across Dataset or Uncertain from Concept:
 - Tattoos or markings if core to the concept
 - Known accessories that always appear (unless outfit-specific)
 
-Updated Caption Format:  
-tr1gger [Style], [Notable Visual Features], [Clothing], [Pose], [Expression], [Lighting], [Camera Angle]
+Caption Format (MUST FOLLOW EXACTLY):  
+tr1gger [Style], [Notable Visual Features], [Clothing], [Pose], [Expression/Mood], [Background/Setting], [Lighting], [Camera Angle]
 
 Captioning Principles:
 - Emphasize visual variation and context-specific details (outfit, pose, lighting, expression, camera angle).
@@ -41,14 +49,20 @@ Captioning Principles:
 - Avoid mentioning real or fictional identities.
 - Always prefix with the trigger word "tr1gger."
 
-Examples:
-- tr1gger photorealistic, tactical vest and gloves, standing in profile, neutral expression, overcast lighting, side profile
-- tr1gger 3D-rendered, hooded cloak with digital pattern, seated cross-legged, calm expression, low ambient lighting, front view
-- tr1gger anime-style, school uniform with blue necktie, standing with arms behind back, gentle smile, soft daylight, three-quarter view
-- tr1gger photorealistic, long trench coat and combat boots, walking through rain-soaked street, determined expression, dramatic shadows, low-angle view
+Examples (MUST FOLLOW THIS EXACT FORMAT):
+tr1gger photorealistic, combat gear, tactical vest and gloves, standing in profile, neutral, empty room, overcast lighting, side profile
+tr1gger 3D-rendered, digital patterns, hooded cloak, seated cross-legged, calm, meditation chamber, low ambient lighting, front view
+tr1gger anime-style, school uniform with blue necktie, standing with arms behind back, gentle smile, classroom, soft daylight, three-quarter view
+tr1gger photorealistic, long trench coat and combat boots, walking, determined, rain-soaked street, dramatic shadows, low-angle view
+
+REMEMBER: Your response must be a single line starting with "tr1gger" and following the exact format above. No additional text, formatting, or explanations are allowed.
 """
 
-def caption_images(images):
+class CaptioningError(Exception):
+    """Exception raised for errors in the captioning process."""
+    pass
+
+def caption_images(images, category=None, batch_mode=False):
     # Convert PIL images to base64 encoded strings
     image_strings = []
     for image in images:
@@ -66,30 +80,103 @@ def caption_images(images):
     client = Together(api_key=api_key)
     captions = []
 
-    # Start a separate chat session for each image
-    for img_str in image_strings:
+    # If batch_mode is True, process all images in a single API call
+    if batch_mode and category:
+        # Create a content array with all images
+        content = [{"type": "text", "text": f"Here is the batch of images for {category}. Please caption each image on a separate line, starting each caption with 'tr1gger'."}]
+        
+        # Add all images to the content array
+        for i, img_str in enumerate(image_strings):
+            content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}})
+            content.append({"type": "text", "text": f"Image {i+1}"})
+        
+        # Send the batch request
         messages = [
             {"role": "system", "content": get_prompt()},
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}},
-                    {"type": "text", "text": "Describe this image."}
-                ]
-            }
+            {"role": "user", "content": content}
         ]
         
-        # Request caption for the image using Llama 4 Maverick
         response = client.chat.completions.create(
             model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
             messages=messages
         )
         
-        # Extract caption from the response
+        # Parse the response to extract captions for each image
         full_response = response.choices[0].message.content.strip()
-        # Post-process to extract only the caption part
-        caption = next((line for line in full_response.splitlines() if line.startswith("tr1gger")), "")
-        captions.append(caption)
+        lines = full_response.splitlines()
+        
+        # Extract captions from the response
+        image_count = len(image_strings)
+        captions = [""] * image_count  # Initialize with empty strings
+        
+        # Extract lines that start with or contain "tr1gger"
+        tr1gger_lines = [line for line in lines if "tr1gger" in line]
+        
+        # Assign captions to images
+        for i in range(image_count):
+            if i < len(tr1gger_lines):
+                caption = tr1gger_lines[i]
+                # If caption contains but doesn't start with tr1gger, extract just that part
+                if not caption.startswith("tr1gger") and "tr1gger" in caption:
+                    caption = caption[caption.index("tr1gger"):]
+                captions[i] = caption
+        
+        # Check if all captions are empty or don't contain the trigger word
+        valid_captions = [c for c in captions if c and "tr1gger" in c]
+        if not valid_captions:
+            error_msg = "Failed to parse any valid captions from batch response. Response contained no lines with 'tr1gger'"
+            error_msg += f"\n\nActual response:\n{full_response}"
+            raise CaptioningError(error_msg)
+        
+        # Check if some captions are missing
+        if len(valid_captions) < len(images):
+            missing_count = len(images) - len(valid_captions)
+            invalid_captions = [(i, c) for i, c in enumerate(captions) if not c or "tr1gger" not in c]
+            error_msg = f"Failed to parse captions for {missing_count} of {len(images)} images in batch mode"
+            error_msg += "\n\nMalformed captions:"
+            for idx, caption in invalid_captions:
+                error_msg += f"\nImage {idx+1}: '{caption}'"
+            raise CaptioningError(error_msg)
+    else:
+        # Original method: process each image separately
+        for img_str in image_strings:
+            messages = [
+                {"role": "system", "content": get_prompt()},
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}},
+                        {"type": "text", "text": "Describe this image."}
+                    ]
+                }
+            ]
+            
+            # Request caption for the image using Llama 4 Maverick
+            response = client.chat.completions.create(
+                model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+                messages=messages
+            )
+            
+            # Extract caption from the response
+            full_response = response.choices[0].message.content.strip()
+            # Post-process to extract only the caption part
+            caption = ""
+            for line in full_response.splitlines():
+                if "tr1gger" in line:
+                    # If caption has a numbered prefix, extract just the caption part
+                    if not line.startswith("tr1gger"):
+                        caption = line[line.index("tr1gger"):]
+                    else:
+                        caption = line
+                    break
+            
+            # Check if caption is valid
+            if not caption:
+                error_msg = "Failed to extract a valid caption (containing 'tr1gger') from the response"
+                error_msg += f"\n\nActual response:\n{full_response}"
+                raise CaptioningError(error_msg)
+                
+            captions.append(caption)
     
     return captions
 
@@ -100,24 +187,3 @@ def extract_captions(file_path):
             if line.startswith("tr1gger"):
                 captions.append(line.strip())
     return captions
-
-# Example usage
-if __name__ == "__main__":
-    if not os.environ.get("TOGETHER_API_KEY"):
-        print("Please update the environment with your Together AI API key.")
-        exit(1)
-
-    # Load images
-    image_paths = ['input/daenyrs_hd.jpg', 'input/girl_body.png']
-    images = [Image.open(path).convert("RGB") for path in image_paths]
-
-    # Generate captions
-    captions = caption_images(images)
-    for i, caption in enumerate(captions):
-        print(f"Generated Caption for Image {i+1}: {caption}")
-
-    # Extract captions from a file
-    file_path = 'post_girl/multiview_0.txt'
-    extracted_captions = extract_captions(file_path)
-    for caption in extracted_captions:
-        print(caption)

demo.py

@@ -0,0 +1,378 @@
+import gradio as gr
+import os
+import zipfile
+from io import BytesIO
+import PIL.Image
+import time
+import tempfile
+from main import process_images, collect_images_by_category, write_captions  # Import the CLI functions
+from dotenv import load_dotenv
+from pathlib import Path
+
+# Load environment variables
+load_dotenv()
+
+# Maximum number of images
+MAX_IMAGES = 30
+
+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 the same code path as CLI"""
+    try:
+        print(f"Processing {len(image_paths)} images, batch_by_category={batch_by_category}")
+        # Create a temporary directory to store the images
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Copy images to temp directory and maintain original order
+            temp_image_paths = []
+            original_to_temp = {}  # Map original paths to temp paths
+            for path in image_paths:
+                filename = os.path.basename(path)
+                temp_path = os.path.join(temp_dir, filename)
+                with open(path, 'rb') as src, open(temp_path, 'wb') as dst:
+                    dst.write(src.read())
+                temp_image_paths.append(temp_path)
+                original_to_temp[path] = temp_path
+            
+            print(f"Created {len(temp_image_paths)} temporary files")
+            
+            # Convert temp_dir to Path object for collect_images_by_category
+            temp_dir_path = Path(temp_dir)
+            
+            # Process images using the CLI code path
+            images_by_category, image_paths_by_category = collect_images_by_category(temp_dir_path)
+            print(f"Collected images into {len(images_by_category)} categories")
+            
+            # Get all images and paths in the correct order
+            all_images = []
+            all_image_paths = []
+            for path in image_paths:  # Use original order
+                temp_path = original_to_temp[path]
+                found = False
+                for category, paths in image_paths_by_category.items():
+                    if temp_path in [str(p) for p in paths]:  # Convert Path objects to strings for comparison
+                        idx = [str(p) for p in paths].index(temp_path)
+                        all_images.append(images_by_category[category][idx])
+                        all_image_paths.append(path)  # Use original path
+                        found = True
+                        break
+                if not found:
+                    print(f"Warning: Could not find image {path} in categorized data")
+            
+            print(f"Collected {len(all_images)} images in correct order")
+            
+            # Process based on batch setting
+            if batch_by_category:
+                # Process each category separately
+                captions = [""] * len(image_paths)  # Initialize with empty strings
+                for category, images in images_by_category.items():
+                    category_paths = image_paths_by_category[category]
+                    print(f"Processing category '{category}' with {len(images)} images")
+                    # Use the same code path as CLI
+                    from caption import caption_images
+                    category_captions = caption_images(images, category=category, batch_mode=True)
+                    print(f"Generated {len(category_captions)} captions for category '{category}'")
+                    print("Category captions:", category_captions)  # Debug print category captions
+                    
+                    # Map captions back to original paths
+                    for temp_path, caption in zip(category_paths, category_captions):
+                        temp_path_str = str(temp_path)
+                        for orig_path, orig_temp in original_to_temp.items():
+                            if orig_temp == temp_path_str:
+                                idx = image_paths.index(orig_path)
+                                captions[idx] = caption
+                                break
+            else:
+                # Process all images at once
+                from caption import caption_images
+                print(f"Processing all {len(all_images)} images at once")
+                all_captions = caption_images(all_images, batch_mode=False)
+                print(f"Generated {len(all_captions)} captions")
+                print("All captions:", all_captions)  # Debug print all captions
+                captions = [""] * len(image_paths)
+                for path, caption in zip(all_image_paths, all_captions):
+                    idx = image_paths.index(path)
+                    captions[idx] = caption
+            
+            print(f"Returning {len(captions)} captions")
+            print("Final captions:", captions)  # Debug print final captions
+            return captions
+            
+    except Exception as e:
+        print(f"Error in processing: {e}")
+        raise
+
+# Main Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Image Autocaptioner")
+    
+    # Store uploaded images
+    stored_image_paths = gr.State([])
+    batch_by_category = gr.State(True)  # State to track if batch by category is enabled
+    
+    # Upload component
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("### Upload your images")
+            image_upload = gr.File(
+                file_count="multiple", 
+                label="Drop your files here", 
+                file_types=["image"],
+                type="filepath"
+            )
+        
+        with gr.Column(scale=1):
+            autocaption_btn = gr.Button("Autocaption Images", variant="primary", interactive=False)
+            status_text = gr.Markdown("Upload images to begin", visible=True)
+            
+            # Advanced settings dropdown
+            with gr.Accordion("Advanced", open=False):
+                batch_category_checkbox = gr.Checkbox(
+                    label="Batch by category", 
+                    value=True,
+                    info="Group similar images together when processing"
+                )
+    
+    # Create a container for the captioning area (initially hidden)
+    with gr.Column(visible=False) as captioning_area:
+        gr.Markdown("### Your images and captions")
+        
+        # Create individual image and caption rows
+        image_rows = []
+        image_components = []
+        caption_components = []
+        
+        for i in range(MAX_IMAGES):
+            with gr.Row(visible=False) as img_row:
+                image_rows.append(img_row)
+                
+                img = gr.Image(
+                    label=f"Image {i+1}",
+                    type="filepath",
+                    show_label=False, 
+                    height=200,
+                    width=200,
+                    scale=1
+                )
+                image_components.append(img)
+                
+                caption = gr.Textbox(
+                    label=f"Caption {i+1}",
+                    lines=3,
+                    scale=2
+                )
+                caption_components.append(caption)
+        
+        # Add download button
+        download_btn = gr.Button("Download Images with Captions", variant="secondary", interactive=False)
+        download_output = gr.File(label="Download Zip", visible=False)
+
+    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),  # autocaption_btn
+            gr.update(interactive=True),  # download_btn
+            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 CLI code path"""
+        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")
+            captions = process_uploaded_images(image_paths, batch_category)
+            print(f"Generated {len(captions)} captions")
+            print("Sample captions:", captions[:2])  # Debug print first two captions
+            
+            gr.Info("Captioning complete!")
+            status = gr.update(value="✅ Captioning complete")
+        except Exception as e:
+            print(f"Error in captioning: {str(e)}")
+            gr.Error(f"Captioning failed: {str(e)}")
+            captions = [f"Error: {str(e)}" for _ in image_paths]
+            status = gr.update(value=f"❌ Error: {str(e)}")
+        
+        # Update caption textboxes
+        caption_updates = []
+        for i in range(MAX_IMAGES):
+            if i < len(captions):
+                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
+    
+    # Update the upload_outputs
+    upload_outputs = [
+        stored_image_paths,
+        captioning_area,
+        autocaption_btn,
+        download_btn,
+        download_output,
+        status_text,
+        *image_rows
+    ]
+    
+    # Update both paths and images in a single flow
+    def process_upload(files):
+        # First get paths and visibility updates
+        image_paths, captioning_update, autocaption_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, autocaption_update, download_btn_update, download_output_update, status_update] + row_updates + image_updates + caption_label_updates
+    
+    # Combined outputs for both functions
+    combined_outputs = upload_outputs + image_components + caption_components
+    
+    image_upload.change(
+        process_upload,
+        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]
+    )
+    
+    # Manage the captioning status
+    def on_captioning_start():
+        return gr.update(value="⏳ Processing captions... please wait"), gr.update(interactive=False)
+    
+    def on_captioning_complete():
+        return gr.update(value="✅ Captioning complete"), gr.update(interactive=True)
+    
+    # Set up captioning button
+    autocaption_btn.click(
+        on_captioning_start,
+        inputs=None,
+        outputs=[status_text, autocaption_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, autocaption_btn]
+    )
+    
+    # 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),
+        inputs=None,
+        outputs=[download_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

main.py

@@ -43,32 +43,16 @@ def validate_input_directory(input_dir):
         sys.exit(1)
     
     if text_files:
-        print("Error: Text files detected in the input directory.")
-        print("The input directory should only contain image files to prevent overwriting existing text files.")
-        print("The following text files were found:")
+        print("Warning: Text files detected in the input directory.")
+        print("The following text files will be overwritten:")
         for file in text_files:
             print(f"  - {file}")
-        sys.exit(1)
 
-def process_images(input_dir, output_dir, fix_outfit=False, batch_images=False):
-    """Process all images in the input directory and generate captions."""
-    input_path = Path(input_dir)
-    output_path = Path(output_dir) if output_dir else input_path
-    
-    # Validate the input directory first
-    validate_input_directory(input_dir)
-    
-    # Create output directory if it doesn't exist
-    os.makedirs(output_path, exist_ok=True)
-    
-    # Track the number of processed images
-    processed_count = 0
-    
-    # Collect all images into a dictionary grouped by category
+def collect_images_by_category(input_path):
+    """Collect all valid images and group them by category."""
     images_by_category = {}
     image_paths_by_category = {}
 
-    # Get all files in the input directory
     for file_path in input_path.iterdir():
         if file_path.is_file() and is_image_file(file_path.name):
             try:
@@ -87,6 +71,53 @@ def process_images(input_dir, output_dir, fix_outfit=False, batch_images=False):
                 image_paths_by_category[category].append(file_path)
             except Exception as e:
                 print(f"Error loading {file_path.name}: {e}")
+    
+    return images_by_category, image_paths_by_category
+
+def process_by_category(images_by_category, image_paths_by_category, input_path, output_path):
+    """Process images in batches by category."""
+    processed_count = 0
+    
+    for category, images in images_by_category.items():
+        image_paths = image_paths_by_category[category]
+        try:
+            # Generate captions for the entire category using batch mode
+            captions = caption_images(images, category=category, batch_mode=True)
+            write_captions(image_paths, captions, input_path, output_path)
+            processed_count += len(images)
+        except Exception as e:
+            print(f"Error generating captions for category '{category}': {e}")
+    
+    return processed_count
+
+def process_all_at_once(images_by_category, image_paths_by_category, input_path, output_path):
+    """Process all images at once."""
+    all_images = [img for imgs in images_by_category.values() for img in imgs]
+    all_image_paths = [path for paths in image_paths_by_category.values() for path in paths]
+    processed_count = 0
+    
+    try:
+        captions = caption_images(all_images, batch_mode=False)
+        write_captions(all_image_paths, captions, input_path, output_path)
+        processed_count += len(all_images)
+    except Exception as e:
+        print(f"Error generating captions: {e}")
+    
+    return processed_count
+
+def process_images(input_dir, output_dir, fix_outfit=False, batch_images=False):
+    """Process all images in the input directory and generate captions."""
+    input_path = Path(input_dir)
+    output_path = Path(output_dir) if output_dir else input_path
+    
+    # Validate the input directory first
+    validate_input_directory(input_dir)
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+    
+    # Collect images by category
+    images_by_category, image_paths_by_category = collect_images_by_category(input_path)
 
     # Log the number of images found
     total_images = sum(len(images) for images in images_by_category.values())
@@ -96,27 +127,11 @@ def process_images(input_dir, output_dir, fix_outfit=False, batch_images=False):
         print("No valid images found to process.")
         return
 
-    # Process images by category if batch_images is True
+    # Process images based on batch setting
     if batch_images:
-        for category, images in images_by_category.items():
-            image_paths = image_paths_by_category[category]
-            try:
-                # Generate captions for the entire category
-                captions = caption_images(images)
-                write_captions(image_paths, captions, input_path, output_path)
-                processed_count += len(images)
-            except Exception as e:
-                print(f"Error generating captions for category '{category}': {e}")
+        processed_count = process_by_category(images_by_category, image_paths_by_category, input_path, output_path)
     else:
-        # Process all images at once if batch_images is False
-        all_images = [img for imgs in images_by_category.values() for img in imgs]
-        all_image_paths = [path for paths in image_paths_by_category.values() for path in paths]
-        try:
-            captions = caption_images(all_images)
-            write_captions(all_image_paths, captions, input_path, output_path)
-            processed_count += len(all_images)
-        except Exception as e:
-            print(f"Error generating captions: {e}")
+        processed_count = process_all_at_once(images_by_category, image_paths_by_category, input_path, output_path)
 
     print(f"\nProcessing complete. {processed_count} images were captioned.")
 

requirements.txt

@@ -0,0 +1,5 @@
+gradio==4.44.1
+Pillow==10.0.0
+pydantic>=2.0.0
+together
+fastapi>=0.100.0