Updated generate_image.py

generate_image.py

@@ -6,19 +6,14 @@ from typing import List, Optional, Tuple
 from runwayml import RunwayML
 import mimetypes
 from urllib.parse import urlparse
-import replicate
 
 def encode_image_to_data_uri(image_path: str) -> str:
     """Convert a local image file to a data URI."""
-    # Get the MIME type
     mime_type, _ = mimetypes.guess_type(image_path)
     if not mime_type or not mime_type.startswith('image/'):
         raise ValueError(f"Unsupported image type for {image_path}")
-    
-    # Read and encode the image
     with open(image_path, 'rb') as image_file:
         encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
-    
     return f"data:{mime_type};base64,{encoded_string}"
 
 def save_generated_image(image_url: str, filename: str = None, batch_folder: str = None) -> str:
@@ -33,33 +28,21 @@ def save_generated_image(image_url: str, filename: str = None, batch_folder: str
     Returns:
         Path to the saved image file
     """
-    # Create batch folder if not provided
     if not batch_folder:
         timestamp = time.strftime("%Y%m%d_%H%M%S")
         batch_folder = f"batch_{timestamp}"
-    
-    # Create directory structure
     output_dir = os.path.join("output", batch_folder)
     os.makedirs(output_dir, exist_ok=True)
-    
-    # Generate filename if not provided
     if not filename:
         timestamp = int(time.time())
         filename = f"generated_{timestamp}.jpg"
-    
-    # Ensure filename has extension
     if not os.path.splitext(filename)[1]:
         filename += ".jpg"
-    
     output_path = os.path.join(output_dir, filename)
-    
-    # Download and save the image
     response = requests.get(image_url)
     response.raise_for_status()
-    
     with open(output_path, 'wb') as f:
         f.write(response.content)
-    
     return output_path
 
 def generate_image_with_references(
@@ -82,37 +65,23 @@ def generate_image_with_references(
         seed: Optional seed for reproducible results
         api_key: Optional API key (uses RUNWAYML_API_SECRET env var if not provided)
         auto_tag_prompt: Whether to automatically append tags to prompt (default: True)
-                        When False, expects user to manually include @character, @scene, @style in prompt
     
     Returns:
         Task ID for the generation request
     """
-    # Initialize client
     client = RunwayML(api_key=api_key or os.environ.get("RUNWAYML_API_SECRET"))
-    
-    # Validate inputs
     if len(reference_image_paths) > 3:
         raise ValueError("Maximum 3 reference images allowed")
-    
     if len(prompt_text) > 1000:
         raise ValueError("Prompt text must be 1000 characters or less")
-    
-    # Prepare reference images with standardized tags
     reference_images = []
     tags = []
-    
-    # Keep track of used standard tags to avoid duplicates
     used_standard_tags = set()
-    
     for i, image_path in enumerate(reference_image_paths):
         if not os.path.exists(image_path):
             raise FileNotFoundError(f"Image file not found: {image_path}")
-        
-        # Create tag based on path structure, prioritizing standard categories
         filename = os.path.splitext(os.path.basename(image_path))[0]
         path_parts = image_path.split(os.sep)
-        
-        # Look for standard category directories
         tag = None
         for part in path_parts:
             if part == 'characters' and 'character' not in used_standard_tags:
@@ -127,66 +96,40 @@ def generate_image_with_references(
                 tag = 'style'
                 used_standard_tags.add('style')
                 break
-        
-        # If no standard category found, create a custom tag from filename
         if not tag:
             tag = f"ref_{filename}".replace('-', '_').replace(' ', '_')[:16]
-            # Ensure tag starts with letter and is alphanumeric + underscore
             tag = ''.join(c for c in tag if c.isalnum() or c == '_')
             if not tag[0].isalpha():
                 tag = f"img_{tag}"
-            tag = tag[:16]  # Ensure max 16 characters
-        
+            tag = tag[:16]
         tags.append(tag)
-        
-        # Convert to data URI
         data_uri = encode_image_to_data_uri(image_path)
-        
-        reference_images.append({
-            "uri": data_uri,
-            "tag": tag
-        })
-    
-    # Handle prompt modification based on auto_tag_prompt setting
+        reference_images.append({"uri": data_uri, "tag": tag})
     final_prompt = prompt_text
     if auto_tag_prompt and tags:
-        # Auto-append tags to prompt
         tag_mentions = " ".join([f"@{tag}" for tag in tags])
         final_prompt = f"{prompt_text} using references: {tag_mentions}"
-        
-        # Ensure we don't exceed character limit
         if len(final_prompt) > 1000:
-            # Try without the descriptive text
             tag_mentions = " ".join([f"@{tag}" for tag in tags])
             final_prompt = f"{prompt_text} {tag_mentions}"
-            
-            # If still too long, truncate prompt text
             if len(final_prompt) > 1000:
                 available_chars = 1000 - len(tag_mentions) - 1
                 final_prompt = f"{prompt_text[:available_chars]} {tag_mentions}"
-    
     print(f"Using tags: {tags}")
     if auto_tag_prompt:
         print(f"Auto-tagged prompt: {final_prompt}")
     else:
         print(f"Manual tagging mode - use @{', @'.join(tags)} in your prompt")
         print(f"Original prompt: {final_prompt}")
-    
-    # Prepare the request parameters
     create_params = {
         "model": model,
         "prompt_text": final_prompt,
         "ratio": ratio,
         "reference_images": reference_images
     }
-    
-    # Only include seed if it's not None
     if seed is not None:
         create_params["seed"] = seed
-    
-    # Create the generation task
     task = client.text_to_image.create(**create_params)
-    
     return task.id
 
 def check_task_status(task_id: str, api_key: Optional[str] = None):
@@ -234,7 +177,6 @@ def generate_and_wait_for_result(
     Returns:
         Tuple of (task_id, saved_image_path)
     """
-    # Start the generation task
     task_id = generate_image_with_references(
         prompt_text=prompt_text,
         reference_image_paths=reference_image_paths,
@@ -244,263 +186,79 @@ def generate_and_wait_for_result(
         api_key=api_key,
         auto_tag_prompt=auto_tag_prompt
     )
-    
     print(f"Image generation started. Task ID: {task_id}")
     print(f"Checking status every {wait_interval} seconds (max {max_retries} attempts)...")
-    
-    # Wait and check status
     for attempt in range(max_retries):
         print(f"Attempt {attempt + 1}/{max_retries} - Waiting {wait_interval} seconds...")
         time.sleep(wait_interval)
-        
         try:
             status = check_task_status(task_id, api_key)
             print(f"Status: {status.status}")
-            
             if status.status == "SUCCEEDED":
                 if hasattr(status, 'output') and status.output:
                     image_url = status.output[0]
                     print(f"Generation completed! Image URL: {image_url}")
-                    
-                    # Save the image
                     saved_path = save_generated_image(image_url, filename, batch_folder)
                     print(f"Image saved to: {saved_path}")
-                    
                     return task_id, saved_path
                 else:
                     print("Task succeeded but no output found")
                     return task_id, None
-            
             elif status.status == "FAILED":
                 print("Task failed")
                 return task_id, None
-            
             elif status.status in ["PENDING", "RUNNING"]:
                 print("Task still in progress...")
                 continue
-            
         except Exception as e:
             print(f"Error checking status: {e}")
             if attempt == max_retries - 1:
                 print("Max retries reached. Task may still be processing.")
                 return task_id, None
-    
     print(f"Timeout after {max_retries} attempts. Task may still be processing.")
     print(f"You can manually check status later using task ID: {task_id}")
     return task_id, None
 
-def generate_image_with_replicate_imagen(
-    prompt: str,
-    aspect_ratio: str = "1:1",
-    output_format: str = "jpg",
-    model: str = "google/imagen-4-fast",
-    safety_filter_level: str = "block_only_high",
-    filename: str = None,
-    api_token: Optional[str] = None
-) -> str:
-    """
-    Generate an image using Replicate's Google Imagen models.
-    
-    Args:
-        prompt: Text prompt for image generation
-        aspect_ratio: Aspect ratio of the generated image (default: "1:1")
-        output_format: Format of the output image (default: "jpg")
-        model: Imagen model to use (default: "google/imagen-4-fast")
-        safety_filter_level: Safety filter level (default: "block_only_high")
-        filename: Optional filename for saved image
-        api_token: Optional API token (uses REPLICATE_API_TOKEN env var if not provided)
-    
-    Returns:
-        Path to the saved image file
-    """
-    # Set API token
-    if api_token:
-        os.environ["REPLICATE_API_TOKEN"] = api_token
-    elif not os.environ.get("REPLICATE_API_TOKEN"):
-        raise ValueError("REPLICATE_API_TOKEN environment variable must be set or api_token must be provided")
-    
-    print(f"Generating image with model: {model}")
-    print(f"Prompt: {prompt}")
-    print(f"Aspect ratio: {aspect_ratio}, Format: {output_format}")
-    
-    # Run the model
-    try:
-        output = replicate.run(
-            model,
-            input={
-                "prompt": prompt,
-                "aspect_ratio": aspect_ratio,
-                "output_format": output_format,
-                "safety_filter_level": safety_filter_level
-            }
-        )
-        
-        # The output is a URL string
-        image_url = output
-        print(f"Image generated successfully: {image_url}")
-        
-        # Save the image
-        saved_path = save_generated_image(image_url, filename)
-        print(f"Image saved to: {saved_path}")
-        
-        return saved_path
-        
-    except Exception as e:
-        print(f"Error generating image with Replicate Imagen: {e}")
-        raise
-
 def main():
-    """Example usage with model selection between runway and imagen-fast."""
-    print("=== Image Generation Model Selection ===")
-    print("Available models:")
-    print("1. runway - RunwayML with reference images")
-    print("2. imagen-fast - Replicate's Google Imagen 4 Fast")
-    
-    model_choice = input("Enter model choice (runway/imagen-fast): ").strip().lower()
-    
-    if model_choice == "runway":
-        print("\n=== Testing RunwayML with Reference Images ===")
-        # Example reference images
-        reference_images = [
-            "assets/characters/japanese_guy.jpg",
-            "assets/scenes/f1-fields.jpg", 
-            "assets/styles/f1-cockpit.jpg"
-        ]
-        
-        print("=== Manual Tagging Mode (Default for Testing) ===")
-        # Example with manual tagging (auto_tag_prompt=False)
-        manual_prompt = "@character in a @scene with @style composition, cinematic lighting, high detail"
-        
-        try:
-            task_id, saved_path = generate_and_wait_for_result(
-                prompt_text=manual_prompt,
-                reference_image_paths=reference_images,
-                ratio="1920:1080",
-                filename="f1_driver_manual_tags.jpg",
-                auto_tag_prompt=False  # Manual tagging mode
-            )
-            
-            if saved_path:
-                print(f"Manual tagging success! Image saved to: {saved_path}")
-            else:
-                print(f"Manual tagging incomplete. Task ID: {task_id}")
-                
-        except Exception as e:
-            print(f"Manual tagging error: {e}")
-        
-        print("\n" + "="*50)
-        print("=== Auto Tagging Mode Example ===")
-        # Example with automatic tagging (auto_tag_prompt=True)
-        auto_prompt = "A Japanese F1 driver in a cockpit style setting on a racing field, cinematic lighting, high detail"
-        
-        try:
-            task_id, saved_path = generate_and_wait_for_result(
-                prompt_text=auto_prompt,
-                reference_image_paths=reference_images,
-                ratio="1920:1080",
-                filename="f1_driver_auto_tags.jpg",
-                auto_tag_prompt=True  # Auto tagging mode
-            )
-            
-            if saved_path:
-                print(f"Auto tagging success! Image saved to: {saved_path}")
-            else:
-                print(f"Auto tagging incomplete. Task ID: {task_id}")
-                
-        except Exception as e:
-            print(f"Auto tagging error: {e}")
-    
-    elif model_choice == "imagen-fast":
-        print("\n=== Testing Replicate's Google Imagen 4 Fast ===")
-        
-        # Get prompt from user or use default
-        prompt = input("Enter image prompt (or press Enter for default): ").strip()
-        if not prompt:
-            prompt = "A cinematic shot of a futuristic sports car racing through a neon-lit cyberpunk city at night, high detail, dramatic lighting"
-        
-        # Get aspect ratio
-        aspect_ratio = input("Enter aspect ratio (default 16:9): ").strip() or "16:9"
-        
-        # Get model version
-        model_version = input("Enter model version (fast/ultra, default fast): ").strip().lower() or "fast"
-        model_name = "google/imagen-4-fast" if model_version == "fast" else "google/imagen-4-ultra"
-        
-        try:
-            saved_path = generate_image_with_replicate_imagen(
-                prompt=prompt,
-                aspect_ratio=aspect_ratio,
-                model=model_name,
-                filename="imagen_test.jpg"
-            )
-            print(f"Imagen generation success! Image saved to: {saved_path}")
-            
-        except Exception as e:
-            print(f"Imagen generation error: {e}")
-    
-    else:
-        print(f"Invalid model choice: {model_choice}")
-        print("Please choose either 'runway' or 'imagen-fast'")
-
-def example_manual_tagging():
-    """
-    Example function demonstrating manual tagging mode.
-    When auto_tag_prompt=False, users must include @character, @scene, @style in their prompts.
-    """
+    print("\n=== Testing RunwayML with Reference Images ===")
     reference_images = [
-        "assets/characters/anime_girl.jpg",
-        "assets/scenes/cyberpunk_city.jpg", 
-        "assets/styles/neon_art.jpg"
+        "assets/characters/japanese_guy.jpg",
+        "assets/scenes/f1-fields.jpg",
+        "assets/styles/f1-cockpit.jpg"
     ]
-    
-    # Manual prompt with explicit tag references
-    prompt_with_tags = """
-    A futuristic @character standing in a @cyberpunk @scene 
-    with @style aesthetic, glowing neon lights, 4k resolution
-    """.strip()
-    
-    print("Manual Tagging Example:")
-    print(f"Prompt: {prompt_with_tags}")
-    
+    print("=== Manual Tagging Mode ===")
+    manual_prompt = "@character in a @scene with @style composition, cinematic lighting, high detail"
     try:
         task_id, saved_path = generate_and_wait_for_result(
-            prompt_text=prompt_with_tags,
+            prompt_text=manual_prompt,
             reference_image_paths=reference_images,
-            auto_tag_prompt=False,  # Disabled - expects manual @tags
-            filename="cyberpunk_manual.jpg"
+            ratio="1920:1080",
+            filename="f1_driver_manual_tags.jpg",
+            auto_tag_prompt=False
         )
-        return task_id, saved_path
+        if saved_path:
+            print(f"Manual tagging success! Image saved to: {saved_path}")
+        else:
+            print(f"Manual tagging incomplete. Task ID: {task_id}")
     except Exception as e:
-        print(f"Error in manual tagging example: {e}")
-        return None, None
-
-def example_auto_tagging():
-    """
-    Example function demonstrating auto tagging mode.
-    When auto_tag_prompt=True, tags are automatically appended to the prompt.
-    """
-    reference_images = [
-        "assets/characters/warrior.jpg",
-        "assets/scenes/medieval_castle.jpg", 
-        "assets/styles/oil_painting.jpg"
-    ]
-    
-    # Simple prompt without tag references
-    simple_prompt = "A brave warrior defending a castle, epic fantasy art"
-    
-    print("Auto Tagging Example:")
-    print(f"Original prompt: {simple_prompt}")
-    
+        print(f"Manual tagging error: {e}")
+    print("\n" + "="*50)
+    print("=== Auto Tagging Mode Example ===")
+    auto_prompt = "A Japanese F1 driver in a cockpit style setting on a racing field, cinematic lighting, high detail"
     try:
         task_id, saved_path = generate_and_wait_for_result(
-            prompt_text=simple_prompt,
+            prompt_text=auto_prompt,
             reference_image_paths=reference_images,
-            auto_tag_prompt=True,  # Enabled - automatically adds @tags
-            filename="fantasy_auto.jpg"
+            ratio="1920:1080",
+            filename="f1_driver_auto_tags.jpg",
+            auto_tag_prompt=True
         )
-        return task_id, saved_path
+        if saved_path:
+            print(f"Auto tagging success! Image saved to: {saved_path}")
+        else:
+            print(f"Auto tagging incomplete. Task ID: {task_id}")
     except Exception as e:
-        print(f"Error in auto tagging example: {e}")
-        return None, None
+        print(f"Auto tagging error: {e}")
 
 if __name__ == "__main__":
     main()