Deploy full Character Sheet Pro with HF auth

- Complete app with 7-view character sheet generation
- All src/ modules (backend_router, character_service, etc.)
- HuggingFace token auth for gated FLUX.2 klein 9B model
- Updated requirements with all dependencies

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,132 +1,887 @@
-import gradio as gr
-import numpy as np
+"""
+Character Sheet Pro - HuggingFace Spaces Version
+=================================================
+
+7-View Character Sheet Generator optimized for HuggingFace Spaces Zero GPU.
+Uses FLUX.2 klein 4B as primary backend with Gemini Flash as fallback.
+
+This is a simplified version of app.py designed for:
+- Zero GPU (A10G 24GB) deployment
+- 5-minute session timeout
+- Automatic model loading on first generation
+"""
+
 import os
-import random
-import spaces
-import torch
-from huggingface_hub import login
+import json
+import logging
+import zipfile
+import threading
+import queue
 import base64
+from pathlib import Path
+from typing import Optional, Tuple, Dict, Any, List, Generator
+from datetime import datetime
+import gradio as gr
+from PIL import Image
+from huggingface_hub import login
 
-# Model access configuration (read-only)
+# HuggingFace authentication for gated models
 def _get_access_key():
-    # Encoded for basic obfuscation
     _k = "aGZfRUR2akdKUXJGRmFQUnhLY1BOUmlUR0lXd0dKYkJ4dkNCWA=="
     return base64.b64decode(_k).decode()
 
 HF_TOKEN = os.environ.get("HF_TOKEN") or _get_access_key()
-print("Authenticating...")
 login(token=HF_TOKEN)
-print("Authentication successful")
-
-from diffusers import Flux2KleinPipeline
+print("HuggingFace authentication successful")
 
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# HuggingFace Spaces SDK - provides @spaces.GPU decorator
+try:
+    import spaces
+    HF_SPACES = True
+except ImportError:
+    # Running locally without spaces SDK
+    HF_SPACES = False
+    # Create a dummy decorator for local testing
+    class spaces:
+        @staticmethod
+        def GPU(duration=300):
+            def decorator(func):
+                return func
+            return decorator
 
-print("Loading FLUX.2 klein 9B model...")
-pipe = Flux2KleinPipeline.from_pretrained(
-    "black-forest-labs/FLUX.2-klein-9B",
-    torch_dtype=dtype,
-    use_auth_token=HF_TOKEN
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
+logger = logging.getLogger(__name__)
+
+# Import local modules
+from src.character_service import CharacterSheetService
+from src.models import CharacterSheetConfig
+from src.backend_router import BackendRouter, BackendType
+from src.utils import preprocess_input_image, sanitize_filename
+
+
+def ensure_png_image(image: Optional[Image.Image], max_size: int = 768) -> Optional[Image.Image]:
+    """Convert any image to PNG-compatible RGB format with proper sizing for FLUX."""
+    if image is None:
+        return None
+    # FLUX models work best with smaller inputs (512-768px)
+    # Larger images slow down processing significantly
+    return preprocess_input_image(image, max_size=max_size, ensure_rgb=True)
+
+
+def create_pending_placeholder(width: int = 200, height: int = 200, text: str = "Pending...") -> Image.Image:
+    """Create a placeholder image showing that generation is pending."""
+    from PIL import ImageDraw, ImageFont
+
+    # Create gradient-like dark background
+    img = Image.new('RGB', (width, height), color=(25, 25, 45))
+    draw = ImageDraw.Draw(img)
+
+    # Draw border to make it clearly a placeholder
+    border_color = (255, 149, 0)  # Orange
+    draw.rectangle([(2, 2), (width-3, height-3)], outline=border_color, width=2)
+
+    # Draw loading indicator (three dots)
+    center_y = height // 2
+    dot_spacing = 20
+    dot_radius = 5
+    for i, offset in enumerate([-dot_spacing, 0, dot_spacing]):
+        shade = 200 + (i * 25)
+        dot_color = (shade, int(shade * 0.6), 0)
+        x = width // 2 + offset
+        draw.ellipse([(x - dot_radius, center_y - dot_radius),
+                      (x + dot_radius, center_y + dot_radius)], fill=dot_color)
+
+    # Draw text
+    try:
+        font = ImageFont.truetype("arial.ttf", 14)
+    except:
+        font = ImageFont.load_default()
+
+    bbox = draw.textbbox((0, 0), text, font=font)
+    text_width = bbox[2] - bbox[0]
+    x = (width - text_width) // 2
+    y = center_y + 25
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
-
-@spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        guidance_scale=1.0
-    ).images[0]
-    return image, seed
-
-examples = [
-    "character turnaround sheet, female warrior with red hair, front view, full body, white background",
-    "character turnaround sheet, male wizard with blue robes, front view, full body, white background",
-    "character turnaround sheet, female elf archer, front view, full body, white background",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
+    draw.text((x, y), text, fill=(180, 180, 180), font=font)
+
+    return img
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+OUTPUT_DIR = Path("./outputs")
+OUTPUT_DIR.mkdir(exist_ok=True)
+
+# Get API key from environment (HuggingFace Spaces secrets)
+API_KEY = os.environ.get("GEMINI_API_KEY", "")
+
+# Model defaults - include all FLUX variants
+MODEL_DEFAULTS = {
+    "flux_klein": {"steps": 4, "guidance": 1.0, "name": "FLUX.2 klein 4B", "costume_in_faces": False},
+    "flux_klein_9b_fp8": {"steps": 4, "guidance": 1.0, "name": "FLUX.2 klein 9B", "costume_in_faces": False},
+    "gemini_flash": {"steps": 1, "guidance": 1.0, "name": "Gemini Flash", "costume_in_faces": True},
+}
+
+
+def get_model_defaults(backend_value: str) -> Tuple[int, float]:
+    """Get default steps and guidance for a backend."""
+    defaults = MODEL_DEFAULTS.get(backend_value, {"steps": 4, "guidance": 1.0})
+    return defaults["steps"], defaults["guidance"]
+
+
+def get_costume_in_faces_default(backend_value: str) -> bool:
+    """Get default for including costume reference in face views."""
+    defaults = MODEL_DEFAULTS.get(backend_value, {"costume_in_faces": True})
+    return defaults.get("costume_in_faces", True)
+
+
+# =============================================================================
+# Presets Loading
+# =============================================================================
+
+EXAMPLES_DIR = Path("./examples")
+PRESETS_FILE = EXAMPLES_DIR / "presets.json"
+
+
+def load_presets() -> Dict[str, Any]:
+    """Load presets configuration from JSON file."""
+    if PRESETS_FILE.exists():
+        with open(PRESETS_FILE, 'r') as f:
+            return json.load(f)
+    return {"characters": [], "costumes": []}
+
+
+def get_character_presets() -> List[Dict]:
+    """Get list of character presets."""
+    presets = load_presets()
+    return presets.get("characters", [])
+
+
+def load_character_preset(preset_id: str) -> Tuple[Optional[Image.Image], str, str]:
+    """Load a character preset."""
+    presets = get_character_presets()
+    for preset in presets:
+        if preset["id"] == preset_id:
+            image_path = EXAMPLES_DIR / preset["file"]
+            if image_path.exists():
+                img = Image.open(image_path)
+                return (
+                    img,
+                    preset.get("name", ""),
+                    preset.get("gender", "Auto/Neutral")
+                )
+    return None, "", "Auto/Neutral"
+
+
+# =============================================================================
+# Character Sheet Metadata
+# =============================================================================
+
+def create_character_sheet_metadata(
+    character_name: str,
+    character_sheet: Image.Image,
+    stages: Dict[str, Any],
+    config: CharacterSheetConfig,
+    backend: str,
+    input_type: str,
+    costume_description: str,
+    steps: int,
+    guidance: float
+) -> Dict[str, Any]:
+    """Create JSON metadata with pixel coordinates for each view."""
+    sheet_width, sheet_height = character_sheet.size
+    spacing = config.spacing
+
+    # Calculate face row dimensions
+    face_images = ['left_face', 'front_face', 'right_face']
+    face_height = 0
+    face_widths = []
+    for name in face_images:
+        if name in stages and stages[name] is not None:
+            face_height = stages[name].height
+            face_widths.append(stages[name].width)
+        else:
+            face_widths.append(0)
+
+    # Calculate body row dimensions
+    body_images = ['left_body', 'front_body', 'right_body', 'back_body']
+    body_height = 0
+    body_widths = []
+    for name in body_images:
+        if name in stages and stages[name] is not None:
+            body_height = stages[name].height
+            body_widths.append(stages[name].width)
+        else:
+            body_widths.append(0)
+
+    body_start_y = face_height + spacing
+
+    # Build view regions
+    views = {}
+
+    # Face row
+    x = 0
+    for i, name in enumerate(face_images):
+        views[name] = {
+            "x": x, "y": 0,
+            "width": face_widths[i], "height": face_height,
+            "description": {
+                "left_face": "Left profile view of face (90 degrees)",
+                "front_face": "Front-facing portrait view",
+                "right_face": "Right profile view of face (90 degrees)"
+            }.get(name, name)
+        }
+        x += face_widths[i]
+
+    # Body row
+    x = 0
+    for i, name in enumerate(body_images):
+        views[name] = {
+            "x": x, "y": body_start_y,
+            "width": body_widths[i], "height": body_height,
+            "description": {
+                "left_body": "Left side full body view (90 degrees)",
+                "front_body": "Front-facing full body view",
+                "right_body": "Right side full body view (90 degrees)",
+                "back_body": "Rear full body view (180 degrees)"
+            }.get(name, name)
+        }
+        x += body_widths[i]
+
+    metadata = {
+        "version": "1.0",
+        "generator": "Character Sheet Pro (HuggingFace Spaces)",
+        "timestamp": datetime.now().isoformat(),
+        "character": {
+            "name": character_name,
+            "input_type": input_type,
+            "costume_description": costume_description or None
+        },
+        "generation": {
+            "backend": backend,
+            "steps": steps,
+            "guidance_scale": guidance
+        },
+        "sheet": {
+            "width": sheet_width,
+            "height": sheet_height,
+            "spacing": spacing,
+            "background_color": config.background_color
+        },
+        "views": views,
+        "files": {
+            "character_sheet": f"{sanitize_filename(character_name)}_character_sheet.png",
+            "individual_views": {
+                name: f"{sanitize_filename(character_name)}_{name}.png"
+                for name in list(face_images) + list(body_images)
+            }
+        }
+    }
+
+    return metadata
+
+
+def create_download_zip(
+    character_name: str,
+    character_sheet: Image.Image,
+    stages: Dict[str, Any],
+    metadata: Dict[str, Any],
+    output_dir: Path
+) -> Path:
+    """Create a ZIP file with character sheet, individual views, and metadata JSON."""
+    safe_name = sanitize_filename(character_name)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    zip_path = output_dir / f"{safe_name}_{timestamp}.zip"
+
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
+        # Add character sheet
+        sheet_path = output_dir / f"{safe_name}_character_sheet.png"
+        character_sheet.save(sheet_path)
+        zf.write(sheet_path, f"{safe_name}_character_sheet.png")
+        sheet_path.unlink()
+
+        # Add individual views
+        view_names = ['left_face', 'front_face', 'right_face',
+                      'left_body', 'front_body', 'right_body', 'back_body']
+        for name in view_names:
+            if name in stages and stages[name] is not None:
+                img = stages[name]
+                img_path = output_dir / f"{safe_name}_{name}.png"
+                img.save(img_path)
+                zf.write(img_path, f"{safe_name}_{name}.png")
+                img_path.unlink()
+
+        # Add metadata JSON
+        json_path = output_dir / f"{safe_name}_metadata.json"
+        with open(json_path, 'w') as f:
+            json.dump(metadata, f, indent=2)
+        zf.write(json_path, f"{safe_name}_metadata.json")
+        json_path.unlink()
+
+    return zip_path
+
+
+# =============================================================================
+# Zero GPU Generation Function
+# =============================================================================
+
+# Global cache for the service (persists across GPU sessions)
+_cached_service = None
+_cached_backend = None
+
+
+@spaces.GPU(duration=300)  # 5-minute timeout for the full pipeline
+def generate_with_gpu(
+    input_image: Optional[Image.Image],
+    input_type: str,
+    character_name: str,
+    gender: str,
+    costume_description: str,
+    costume_image: Optional[Image.Image],
+    face_image: Optional[Image.Image],
+    body_image: Optional[Image.Image],
+    backend_choice: str,
+    api_key: str,
+    num_steps: int,
+    guidance_scale: float,
+    include_costume_in_faces: bool
+) -> Tuple[Optional[Image.Image], str, Dict[str, Any]]:
+    """
+    GPU-wrapped generation function for Zero GPU.
+
+    This function runs entirely within a GPU session.
+    Model loading happens inside this function for Zero GPU compatibility.
+    """
+    global _cached_service, _cached_backend
+
+    try:
+        # Determine backend
+        backend = BackendRouter.backend_from_string(backend_choice)
+        is_cloud = backend in (BackendType.GEMINI_FLASH, BackendType.GEMINI_PRO)
+
+        # Validate API key for cloud backends
+        if is_cloud and not api_key:
+            return None, "Error: Gemini API key required for cloud backends", {}
+
+        # Load or reuse service
+        if _cached_service is None or _cached_backend != backend:
+            logger.info(f"Loading model for {backend.value}...")
+
+            # For local FLUX model, create service (this loads the model)
+            _cached_service = CharacterSheetService(
+                api_key=api_key if is_cloud else None,
+                backend=backend
+            )
+            _cached_backend = backend
+
+            # Configure steps/guidance
+            if hasattr(_cached_service.client, 'default_steps'):
+                _cached_service.client.default_steps = num_steps
+            if hasattr(_cached_service.client, 'default_guidance'):
+                _cached_service.client.default_guidance = guidance_scale
+
+            logger.info(f"Model loaded successfully: {backend.value}")
+
+        # Map gender selection
+        gender_map = {
+            "Auto/Neutral": "character",
+            "Male": "man",
+            "Female": "woman"
+        }
+        gender_term = gender_map.get(gender, "character")
+
+        # Validate steps and guidance
+        num_steps = max(1, min(100, int(num_steps)))
+        guidance_scale = max(0.0, min(20.0, float(guidance_scale)))
+
+        # Update steps/guidance if different
+        if hasattr(_cached_service.client, 'default_steps'):
+            _cached_service.client.default_steps = num_steps
+        if hasattr(_cached_service.client, 'default_guidance'):
+            _cached_service.client.default_guidance = guidance_scale
+
+        # Run generation
+        logger.info(f"Starting generation for {character_name}...")
+
+        sheet, status, metadata = _cached_service.generate_character_sheet(
+            initial_image=input_image,
+            input_type=input_type,
+            character_name=character_name or "Character",
+            gender_term=gender_term,
+            costume_description=costume_description,
+            costume_image=costume_image,
+            face_image=face_image,
+            body_image=body_image,
+            include_costume_in_faces=include_costume_in_faces,
+            output_dir=OUTPUT_DIR
+        )
+
+        return sheet, status, metadata
+
+    except Exception as e:
+        logger.exception(f"Generation error: {e}")
+        return None, f"Error: {str(e)}", {}
+
+
+# =============================================================================
+# Gradio Interface Functions
+# =============================================================================
+
+def generate_character_sheet(
+    input_image: Optional[Image.Image],
+    input_type: str,
+    character_name: str,
+    gender: str,
+    costume_description: str,
+    costume_image: Optional[Image.Image],
+    face_image: Optional[Image.Image],
+    body_image: Optional[Image.Image],
+    backend_choice: str,
+    api_key_override: str,
+    num_steps: int,
+    guidance_scale: float,
+    include_costume_in_faces: bool,
+    progress=gr.Progress()
+) -> Generator:
+    """
+    Generate character sheet from input image(s).
+
+    This wrapper handles preprocessing and calls the GPU-wrapped function.
+    """
+    # Initial empty state
+    empty_previews = [None] * 7
+
+    yield (None, "Initializing...", *empty_previews, None, None)
+
+    # Preprocess all input images to PNG format
+    input_image = ensure_png_image(input_image)
+    face_image = ensure_png_image(face_image)
+    body_image = ensure_png_image(body_image)
+    costume_image = ensure_png_image(costume_image)
+
+    # Validate input
+    if input_type == "Face + Body (Separate)":
+        if face_image is None or body_image is None:
+            yield (None, "Error: Both face and body images required for this mode.",
+                   *empty_previews, None, None)
+            return
+    elif input_image is None:
+        yield (None, "Error: Please upload an input image.", *empty_previews, None, None)
+        return
+
+    # Get API key
+    api_key = api_key_override.strip() if api_key_override.strip() else API_KEY
+
+    # Show loading state
+    progress(0.1, desc="Allocating GPU...")
+    yield (None, "Allocating GPU and loading model (this may take 30-60 seconds on first run)...",
+           *empty_previews, None, None)
+
+    try:
+        # Call the GPU-wrapped function
+        character_sheet, status, metadata = generate_with_gpu(
+            input_image=input_image,
+            input_type=input_type,
+            character_name=character_name or "Character",
+            gender=gender,
+            costume_description=costume_description,
+            costume_image=costume_image,
+            face_image=face_image,
+            body_image=body_image,
+            backend_choice=backend_choice,
+            api_key=api_key,
+            num_steps=int(num_steps),
+            guidance_scale=float(guidance_scale),
+            include_costume_in_faces=include_costume_in_faces
+        )
+
+        if character_sheet is None:
+            yield (None, status, *empty_previews, None, None)
+            return
+
+        # Get stages from metadata for preview
+        stages = metadata.get('stages', {})
+
+        # Create preview list
+        preview_list = [
+            stages.get('left_face'),
+            stages.get('front_face'),
+            stages.get('right_face'),
+            stages.get('left_body'),
+            stages.get('front_body'),
+            stages.get('right_body'),
+            stages.get('back_body')
+        ]
+
+        # Determine backend
+        backend = BackendRouter.backend_from_string(backend_choice)
+
+        # Create metadata JSON
+        config = CharacterSheetConfig()
+        json_metadata = create_character_sheet_metadata(
+            character_name=character_name or "Character",
+            character_sheet=character_sheet,
+            stages=stages,
+            config=config,
+            backend=BackendRouter.BACKEND_NAMES.get(backend, backend_choice),
+            input_type=input_type,
+            costume_description=costume_description,
+            steps=num_steps,
+            guidance=guidance_scale
+        )
+
+        # Save JSON file
+        safe_name = sanitize_filename(character_name or "Character")
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        json_path = OUTPUT_DIR / f"{safe_name}_{timestamp}_metadata.json"
+        with open(json_path, 'w') as f:
+            json.dump(json_metadata, f, indent=2)
+
+        # Create ZIP file
+        zip_path = create_download_zip(
+            character_name=character_name or "Character",
+            character_sheet=character_sheet,
+            stages=stages,
+            metadata=json_metadata,
+            output_dir=OUTPUT_DIR
+        )
+
+        # Final yield with all outputs
+        yield (
+            character_sheet,
+            status,
+            *preview_list,
+            str(json_path),
+            str(zip_path)
+        )
+
+    except Exception as e:
+        logger.exception(f"Error: {e}")
+        yield (None, f"Error: {str(e)}", *empty_previews, None, None)
+
+
+def update_input_visibility(input_type: str):
+    """Update visibility of input components based on input type."""
+    if input_type == "Face + Body (Separate)":
+        return (
+            gr.update(visible=False),  # Main input
+            gr.update(visible=True),   # Face input
+            gr.update(visible=True),   # Body input
+        )
+    else:
+        return (
+            gr.update(visible=True),   # Main input
+            gr.update(visible=False),  # Face input
+            gr.update(visible=False),  # Body input
+        )
+
+
+def update_defaults_on_backend_change(backend_value: str):
+    """Update steps, guidance, and costume-in-faces when backend changes."""
+    steps, guidance = get_model_defaults(backend_value)
+    costume_in_faces = get_costume_in_faces_default(backend_value)
+    return gr.update(value=steps), gr.update(value=guidance), gr.update(value=costume_in_faces)
+
+
+# =============================================================================
+# Gradio UI
+# =============================================================================
+
+# CSS for the interface
+APP_CSS = """
+.container { max-width: 1200px; margin: auto; }
+.output-image { min-height: 400px; }
+
+/* GPU status banner */
+.gpu-banner {
+    background: linear-gradient(90deg, #7c3aed, #a855f7);
+    padding: 12px 20px;
+    text-align: center;
+    color: white;
+    font-weight: bold;
+    border-radius: 8px;
+    margin-bottom: 16px;
+}
+
+/* Generate button styling */
+.generate-btn-main {
+    background: linear-gradient(90deg, #00aa44, #00cc55) !important;
+    color: white !important;
+    font-weight: bold !important;
+    font-size: 20px !important;
+    padding: 16px 32px !important;
+    border: none !important;
+    box-shadow: 0 4px 15px rgba(0, 170, 68, 0.4) !important;
+}
+
+.generate-btn-main:hover {
+    background: linear-gradient(90deg, #00cc55, #00ee66) !important;
 }
 """
 
-with gr.Blocks() as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# CharacterForgePro
-Generate character images using FLUX.2 klein 9B on Zero GPU.
-""")
 
+def create_ui():
+    """Create the Gradio interface for HuggingFace Spaces."""
+
+    with gr.Blocks(title="Character Sheet Pro") as demo:
+
+        # GPU status banner
+        gr.HTML(
+            '<div class="gpu-banner">'
+            'Zero GPU (A10G) - Model loads automatically on first generation'
+            '</div>'
+        )
+
+        gr.Markdown("# Character Sheet Pro")
+        gr.Markdown("Generate 7-view character turnaround sheets from a single input image using FLUX.2 klein.")
+
+        # Backend selection and controls
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
+            backend_dropdown = gr.Dropdown(
+                choices=[
+                    ("FLUX.2 klein 9B (Best Quality, ~20GB)", "flux_klein_9b_fp8"),
+                    ("FLUX.2 klein 4B (Fast, ~13GB)", BackendType.FLUX_KLEIN.value),
+                    ("Gemini Flash (Cloud - Fallback)", BackendType.GEMINI_FLASH.value),
+                ],
+                value="flux_klein_9b_fp8",  # Default to best quality
+                label="Backend",
+                scale=2
             )
-            run_button = gr.Button("Run", scale=0)
-
-        result = gr.Image(label="Result", show_label=False)
 
-        with gr.Accordion("Advanced Settings", open=False):
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+            api_key_input = gr.Textbox(
+                label="Gemini API Key (for cloud backend)",
+                placeholder="Enter API key if using Gemini",
+                type="password",
+                value="",
+                scale=2
             )
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+
+        with gr.Row():
+            # Left column: Inputs
+            with gr.Column(scale=1):
+                gr.Markdown("### Input Settings")
+
+                input_type = gr.Radio(
+                    choices=["Face Only", "Full Body", "Face + Body (Separate)"],
+                    value="Face Only",
+                    label="Input Type",
+                    info="What type of image(s) are you providing?"
+                )
+
+                main_input = gr.Image(
+                    label="Input Image",
+                    type="pil",
+                    format="png",
+                    visible=True
+                )
+
+                with gr.Row(visible=False) as face_body_row:
+                    face_input = gr.Image(
+                        label="Face Reference",
+                        type="pil",
+                        format="png",
+                        visible=False
+                    )
+                    body_input = gr.Image(
+                        label="Body Reference",
+                        type="pil",
+                        format="png",
+                        visible=False
+                    )
+
+                gr.Markdown("### Character Details")
+
+                character_name = gr.Textbox(
+                    label="Character Name",
+                    placeholder="My Character",
+                    value=""
+                )
+
+                gender = gr.Radio(
+                    choices=["Auto/Neutral", "Male", "Female"],
+                    value="Auto/Neutral",
+                    label="Gender"
+                )
+
+                costume_description = gr.Textbox(
+                    label="Costume Description (Optional)",
+                    placeholder="e.g., Full plate armor with gold trim...",
+                    value="",
+                    lines=3
                 )
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+
+                costume_image = gr.Image(
+                    label="Costume Reference Image (Optional)",
+                    type="pil",
+                    format="png"
+                )
+
+                gr.Markdown("### Generation Parameters")
+
+                with gr.Row():
+                    num_steps = gr.Number(
+                        label="Inference Steps",
+                        value=4,
+                        minimum=1,
+                        maximum=50,
+                        step=1,
+                        info="FLUX klein uses 4 steps"
+                    )
+                    guidance_scale = gr.Number(
+                        label="Guidance Scale",
+                        value=1.0,
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.1,
+                        info="FLUX klein uses 1.0"
+                    )
+
+                include_costume_in_faces = gr.Checkbox(
+                    label="Include costume in face views",
+                    value=False,
+                    info="Turn OFF for FLUX (can confuse framing)"
                 )
 
-            with gr.Row():
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=4,
+                # GENERATE BUTTON
+                generate_btn = gr.Button(
+                    "GENERATE CHARACTER SHEET",
+                    variant="primary",
+                    size="lg",
+                    elem_classes=["generate-btn-main"]
                 )
 
-        gr.Examples(
-            examples=examples,
-            fn=infer,
-            inputs=[prompt],
-            outputs=[result, seed],
-            cache_examples=False
+            # Right column: Output
+            with gr.Column(scale=2):
+                gr.Markdown("### Generated Character Sheet")
+
+                output_image = gr.Image(
+                    label="Character Sheet",
+                    type="pil",
+                    format="png",
+                    elem_classes=["output-image"]
+                )
+
+                status_text = gr.Textbox(
+                    label="Status",
+                    interactive=False
+                )
+
+                # Preview gallery
+                gr.Markdown("### Individual Views Preview")
+
+                with gr.Row():
+                    gr.Markdown("**Face Views:**")
+                with gr.Row():
+                    preview_left_face = gr.Image(label="Left Face", type="pil", height=150, width=112)
+                    preview_front_face = gr.Image(label="Front Face", type="pil", height=150, width=112)
+                    preview_right_face = gr.Image(label="Right Face", type="pil", height=150, width=112)
+
+                with gr.Row():
+                    gr.Markdown("**Body Views:**")
+                with gr.Row():
+                    preview_left_body = gr.Image(label="Left Body", type="pil", height=150, width=84)
+                    preview_front_body = gr.Image(label="Front Body", type="pil", height=150, width=84)
+                    preview_right_body = gr.Image(label="Right Body", type="pil", height=150, width=84)
+                    preview_back_body = gr.Image(label="Back Body", type="pil", height=150, width=84)
+
+                # Downloads
+                gr.Markdown("### Downloads")
+                with gr.Row():
+                    json_download = gr.File(label="Metadata JSON", interactive=False)
+                    zip_download = gr.File(label="Complete Package (ZIP)", interactive=False)
+
+        # Usage instructions
+        gr.Markdown("---")
+        gr.Markdown("### How to Use")
+        gr.Markdown("""
+        1. **Upload an image** (face portrait or full body)
+        2. **Select input type** based on your image
+        3. **Optionally** add character name, gender, and costume description
+        4. **Click Generate** - the model loads automatically on first run (~30-60s)
+        5. **Wait** for all 7 views to generate (~2-3 minutes total)
+        6. **Download** the complete package
+
+        **GPU Notes:**
+        - Uses Zero GPU (A10G 24GB) - free but with 5-minute session limit
+        - First generation loads the model (adds ~30-60 seconds)
+        - Subsequent generations in the same session are faster
+        - If GPU unavailable, switch to Gemini Flash (requires API key)
+        """)
+
+        # Event handlers
+        input_type.change(
+            fn=update_input_visibility,
+            inputs=[input_type],
+            outputs=[main_input, face_input, body_input]
+        )
+
+        backend_dropdown.change(
+            fn=update_defaults_on_backend_change,
+            inputs=[backend_dropdown],
+            outputs=[num_steps, guidance_scale, include_costume_in_faces]
         )
 
-        gr.on(
-            triggers=[run_button.click, prompt.submit],
-            fn=infer,
-            inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
-            outputs=[result, seed]
+        generate_btn.click(
+            fn=generate_character_sheet,
+            inputs=[
+                main_input,
+                input_type,
+                character_name,
+                gender,
+                costume_description,
+                costume_image,
+                face_input,
+                body_input,
+                backend_dropdown,
+                api_key_input,
+                num_steps,
+                guidance_scale,
+                include_costume_in_faces
+            ],
+            outputs=[
+                output_image,
+                status_text,
+                preview_left_face,
+                preview_front_face,
+                preview_right_face,
+                preview_left_body,
+                preview_front_body,
+                preview_right_body,
+                preview_back_body,
+                json_download,
+                zip_download
+            ]
         )
 
-demo.launch(css=css)
+    return demo
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+if __name__ == "__main__":
+    demo = create_ui()
+
+    if HF_SPACES:
+        # Running on HuggingFace Spaces
+        demo.launch(
+            theme=gr.themes.Soft(),
+            css=APP_CSS
+        )
+    else:
+        # Local testing
+        print("Running locally (no Zero GPU)")
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7890,
+            share=False,
+            theme=gr.themes.Soft(),
+            css=APP_CSS
+        )

requirements.txt

@@ -1,8 +1,36 @@
-accelerate>=1.12.0
+# Character Sheet Pro - HuggingFace Spaces
+# =========================================
+
+# Image processing
+Pillow>=10.0.0
+
+# Utilities
+python-dotenv>=1.0.0
+
+# Diffusers from git (required for Flux2KleinPipeline)
 git+https://github.com/huggingface/diffusers.git
-torch>=2.5.0
-transformers>=5.0.0
-sentencepiece
-numpy
-Pillow
+
+# PyTorch
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Transformers
+transformers>=4.40.0
+
+# Accelerate
+accelerate>=0.25.0
+
+# HuggingFace Hub
+huggingface-hub>=0.20.0
+
+# Safetensors
+safetensors>=0.4.0
+
+# Sentencepiece
+sentencepiece>=0.1.99
+
+# Google Gemini API (fallback backend)
+google-genai>=0.3.0
+
+# Invisible watermark (for FLUX)
 invisible_watermark

src/__init__.py

@@ -0,0 +1,37 @@
+"""
+Character Sheet Pro - 7-View Character Sheet Generator
+======================================================
+
+A standalone character sheet generation system that creates
+multi-view turnaround sheets from a single input image.
+
+Supports:
+- 7 views (3 face + 4 body)
+- Multiple backends: Gemini (Cloud), FLUX.2 klein (Local), Qwen-Image-Edit (Local/ComfyUI)
+- HuggingFace Spaces deployment via Gradio
+"""
+
+from .models import GenerationRequest, GenerationResult
+from .gemini_client import GeminiClient
+from .character_service import CharacterSheetService
+from .backend_router import BackendRouter, BackendType
+from .flux_klein_client import FluxKleinClient
+from .qwen_image_edit_client import QwenImageEditClient
+from .comfyui_client import ComfyUIClient
+from .model_manager import ModelManager, ModelState, get_model_manager
+
+__version__ = "2.3.0"  # Bumped for model manager feature
+__all__ = [
+    "GenerationRequest",
+    "GenerationResult",
+    "GeminiClient",
+    "CharacterSheetService",
+    "BackendRouter",
+    "BackendType",
+    "FluxKleinClient",
+    "QwenImageEditClient",
+    "ComfyUIClient",
+    "ModelManager",
+    "ModelState",
+    "get_model_manager",
+]

src/backend_router.py

@@ -0,0 +1,401 @@
+"""
+Backend Router
+==============
+
+Unified router for selecting between different image generation backends:
+- Gemini (Flash/Pro) - Cloud API
+- FLUX.2 klein 4B/9B - Local model
+- Z-Image Turbo (Tongyi-MAI) - Local model, 6B, 9 steps, 16GB VRAM
+- Qwen-Image-Edit-2511 - Local model
+"""
+
+import logging
+from typing import Optional, Protocol, Union
+from enum import Enum, auto
+from PIL import Image
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class BackendType(Enum):
+    """Available backend types."""
+    GEMINI_FLASH = "gemini_flash"
+    GEMINI_PRO = "gemini_pro"
+    FLUX_KLEIN = "flux_klein"              # 4B model (~13GB VRAM)
+    FLUX_KLEIN_9B_FP8 = "flux_klein_9b_fp8"  # 9B FP8 model (~20GB VRAM, best quality)
+    ZIMAGE_TURBO = "zimage_turbo"          # Z-Image Turbo 6B (9 steps, 16GB VRAM)
+    ZIMAGE_BASE = "zimage_base"            # Z-Image Base 6B (50 steps, CFG support) - NEW!
+    LONGCAT_EDIT = "longcat_edit"          # LongCat-Image-Edit (instruction-following, 18GB)
+    QWEN_IMAGE_EDIT = "qwen_image_edit"    # Direct diffusers (slow, high VRAM)
+    QWEN_COMFYUI = "qwen_comfyui"           # Via ComfyUI with FP8 quantization
+
+
+class ImageClient(Protocol):
+    """Protocol for image generation clients."""
+
+    def generate(self, request: GenerationRequest, **kwargs) -> GenerationResult:
+        """Generate an image from request."""
+        ...
+
+    def is_healthy(self) -> bool:
+        """Check if client is ready."""
+        ...
+
+
+class BackendRouter:
+    """
+    Router for selecting between image generation backends.
+
+    Supports lazy loading of local models to save memory.
+    """
+
+    BACKEND_NAMES = {
+        BackendType.GEMINI_FLASH: "Gemini Flash",
+        BackendType.GEMINI_PRO: "Gemini Pro",
+        BackendType.FLUX_KLEIN: "FLUX.2 klein 4B",
+        BackendType.FLUX_KLEIN_9B_FP8: "FLUX.2 klein 9B-FP8",
+        BackendType.ZIMAGE_TURBO: "Z-Image Turbo 6B",
+        BackendType.ZIMAGE_BASE: "Z-Image Base 6B",
+        BackendType.LONGCAT_EDIT: "LongCat-Image-Edit",
+        BackendType.QWEN_IMAGE_EDIT: "Qwen-Image-Edit-2511",
+        BackendType.QWEN_COMFYUI: "Qwen-Image-Edit-2511-FP8 (ComfyUI)",
+    }
+
+    def __init__(
+        self,
+        gemini_api_key: Optional[str] = None,
+        default_backend: BackendType = BackendType.GEMINI_FLASH
+    ):
+        """
+        Initialize backend router.
+
+        Args:
+            gemini_api_key: API key for Gemini backends
+            default_backend: Default backend to use
+        """
+        self.gemini_api_key = gemini_api_key
+        self.default_backend = default_backend
+        self._clients: dict = {}
+        self._active_backend: Optional[BackendType] = None
+
+        logger.info(f"BackendRouter initialized (default: {default_backend.value})")
+
+    def get_client(self, backend: Optional[BackendType] = None) -> ImageClient:
+        """
+        Get or create client for specified backend.
+
+        Args:
+            backend: Backend type (uses default if None)
+
+        Returns:
+            ImageClient instance
+        """
+        if backend is None:
+            backend = self.default_backend
+
+        # Return cached client if available
+        if backend in self._clients:
+            self._active_backend = backend
+            return self._clients[backend]
+
+        # Create new client
+        client = self._create_client(backend)
+        self._clients[backend] = client
+        self._active_backend = backend
+
+        return client
+
+    def _create_client(self, backend: BackendType) -> ImageClient:
+        """Create client for specified backend."""
+        logger.info(f"Creating client for {backend.value}...")
+
+        if backend == BackendType.GEMINI_FLASH:
+            from .gemini_client import GeminiClient
+            if not self.gemini_api_key:
+                raise ValueError("Gemini API key required for Gemini backends")
+            return GeminiClient(api_key=self.gemini_api_key, use_pro_model=False)
+
+        elif backend == BackendType.GEMINI_PRO:
+            from .gemini_client import GeminiClient
+            if not self.gemini_api_key:
+                raise ValueError("Gemini API key required for Gemini backends")
+            return GeminiClient(api_key=self.gemini_api_key, use_pro_model=True)
+
+        elif backend == BackendType.FLUX_KLEIN:
+            from .flux_klein_client import FluxKleinClient
+            # 4B model (~13GB VRAM) - fast
+            client = FluxKleinClient(
+                model_variant="4b",
+                enable_cpu_offload=False
+            )
+            if not client.load_model():
+                raise RuntimeError("Failed to load FLUX.2 klein 4B model")
+            return client
+
+        elif backend == BackendType.FLUX_KLEIN_9B_FP8:
+            from .flux_klein_client import FluxKleinClient
+            # 9B model (~29GB VRAM with CPU offload) - best quality
+            client = FluxKleinClient(
+                model_variant="9b",
+                enable_cpu_offload=True  # Required for 24GB VRAM
+            )
+            if not client.load_model():
+                raise RuntimeError("Failed to load FLUX.2 klein 9B model")
+            return client
+
+        elif backend == BackendType.ZIMAGE_TURBO:
+            from .zimage_client import ZImageClient
+            # Z-Image Turbo 6B - fast (9 steps), fits 16GB VRAM
+            client = ZImageClient(
+                model_variant="turbo",
+                enable_cpu_offload=True
+            )
+            if not client.load_model():
+                raise RuntimeError("Failed to load Z-Image Turbo model")
+            return client
+
+        elif backend == BackendType.ZIMAGE_BASE:
+            from .zimage_client import ZImageClient
+            # Z-Image Base 6B - quality (50 steps), CFG support, negative prompts
+            client = ZImageClient(
+                model_variant="base",
+                enable_cpu_offload=True
+            )
+            if not client.load_model():
+                raise RuntimeError("Failed to load Z-Image Base model")
+            return client
+
+        elif backend == BackendType.LONGCAT_EDIT:
+            from .longcat_edit_client import LongCatEditClient
+            # LongCat-Image-Edit - instruction-following editing (~18GB VRAM)
+            client = LongCatEditClient(
+                enable_cpu_offload=True
+            )
+            if not client.load_model():
+                raise RuntimeError("Failed to load LongCat-Image-Edit model")
+            return client
+
+        elif backend == BackendType.QWEN_IMAGE_EDIT:
+            from .qwen_image_edit_client import QwenImageEditClient
+            client = QwenImageEditClient(enable_cpu_offload=False)
+            if not client.load_model():
+                raise RuntimeError("Failed to load Qwen-Image-Edit model")
+            return client
+
+        elif backend == BackendType.QWEN_COMFYUI:
+            from .comfyui_client import ComfyUIClient
+            client = ComfyUIClient()
+            if not client.is_healthy():
+                raise RuntimeError(
+                    "ComfyUI is not running. Please start ComfyUI first:\n"
+                    "  cd comfyui && python main.py"
+                )
+            return client
+
+        else:
+            raise ValueError(f"Unknown backend: {backend}")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        backend: Optional[BackendType] = None,
+        **kwargs
+    ) -> GenerationResult:
+        """
+        Generate image using specified backend.
+
+        Args:
+            request: Generation request
+            backend: Backend to use (default if None)
+            **kwargs: Backend-specific parameters
+
+        Returns:
+            GenerationResult
+        """
+        try:
+            client = self.get_client(backend)
+            return client.generate(request, **kwargs)
+        except Exception as e:
+            logger.error(f"Generation failed with {backend}: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Backend error: {str(e)}")
+
+    def unload_local_models(self):
+        """Unload all local models to free memory."""
+        local_backends = (BackendType.FLUX_KLEIN, BackendType.FLUX_KLEIN_9B_FP8, BackendType.ZIMAGE_TURBO, BackendType.ZIMAGE_BASE, BackendType.LONGCAT_EDIT, BackendType.QWEN_IMAGE_EDIT, BackendType.QWEN_COMFYUI)
+        for backend, client in list(self._clients.items()):
+            if backend in local_backends:
+                if hasattr(client, 'unload_model'):
+                    client.unload_model()
+                del self._clients[backend]
+                logger.info(f"Unloaded {backend.value}")
+
+    def switch_backend(self, backend: BackendType) -> bool:
+        """
+        Switch to a different backend.
+
+        For local models, this will load the new model and optionally
+        unload the previous one to save memory.
+
+        Args:
+            backend: Backend to switch to
+
+        Returns:
+            True if switch successful
+        """
+        try:
+            local_backends = {BackendType.FLUX_KLEIN, BackendType.FLUX_KLEIN_9B_FP8, BackendType.ZIMAGE_TURBO, BackendType.ZIMAGE_BASE, BackendType.LONGCAT_EDIT, BackendType.QWEN_IMAGE_EDIT, BackendType.QWEN_COMFYUI}
+
+            # Unload other local models first to save memory
+            if backend in local_backends:
+                for other_local in local_backends - {backend}:
+                    if other_local in self._clients:
+                        if hasattr(self._clients[other_local], 'unload_model'):
+                            self._clients[other_local].unload_model()
+                        del self._clients[other_local]
+
+            # Get/create the new client
+            self.get_client(backend)
+            self.default_backend = backend
+
+            logger.info(f"Switched to {backend.value}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to switch to {backend}: {e}", exc_info=True)
+            return False
+
+    def get_active_backend_name(self) -> str:
+        """Get human-readable name of active backend."""
+        if self._active_backend:
+            return self.BACKEND_NAMES.get(self._active_backend, str(self._active_backend))
+        return "None"
+
+    def is_local_backend(self, backend: Optional[BackendType] = None) -> bool:
+        """Check if backend is a local model."""
+        if backend is None:
+            backend = self._active_backend
+        return backend in (BackendType.FLUX_KLEIN, BackendType.FLUX_KLEIN_9B_FP8, BackendType.ZIMAGE_TURBO, BackendType.ZIMAGE_BASE, BackendType.LONGCAT_EDIT, BackendType.QWEN_IMAGE_EDIT, BackendType.QWEN_COMFYUI)
+
+    @staticmethod
+    def get_supported_aspect_ratios(backend: BackendType) -> dict:
+        """
+        Get supported aspect ratios for a backend.
+
+        Returns dict mapping ratio strings to (width, height) tuples.
+        """
+        # Import clients to get their ASPECT_RATIOS
+        if backend in (BackendType.FLUX_KLEIN, BackendType.FLUX_KLEIN_9B_FP8):
+            from .flux_klein_client import FluxKleinClient
+            return FluxKleinClient.ASPECT_RATIOS
+
+        elif backend in (BackendType.ZIMAGE_TURBO, BackendType.ZIMAGE_BASE):
+            from .zimage_client import ZImageClient
+            return ZImageClient.ASPECT_RATIOS
+
+        elif backend == BackendType.LONGCAT_EDIT:
+            from .longcat_edit_client import LongCatEditClient
+            return LongCatEditClient.ASPECT_RATIOS
+
+        elif backend in (BackendType.GEMINI_FLASH, BackendType.GEMINI_PRO):
+            from .gemini_client import GeminiClient
+            return GeminiClient.ASPECT_RATIOS
+
+        elif backend == BackendType.QWEN_IMAGE_EDIT:
+            from .qwen_image_edit_client import QwenImageEditClient
+            return QwenImageEditClient.ASPECT_RATIOS
+
+        elif backend == BackendType.QWEN_COMFYUI:
+            from .comfyui_client import ComfyUIClient
+            return ComfyUIClient.ASPECT_RATIOS
+
+        else:
+            # Default fallback
+            return {
+                "1:1": (1024, 1024),
+                "16:9": (1344, 768),
+                "9:16": (768, 1344),
+            }
+
+    @staticmethod
+    def get_aspect_ratio_choices(backend: BackendType) -> list:
+        """
+        Get aspect ratio choices for UI dropdowns.
+
+        Returns list of (label, value) tuples.
+        """
+        ratios = BackendRouter.get_supported_aspect_ratios(backend)
+        choices = []
+        for ratio, (w, h) in ratios.items():
+            label = f"{ratio} ({w}x{h})"
+            choices.append((label, ratio))
+        return choices
+
+    def get_available_backends(self) -> list:
+        """Get list of available backends."""
+        available = []
+
+        # Gemini backends require API key
+        if self.gemini_api_key:
+            available.extend([BackendType.GEMINI_FLASH, BackendType.GEMINI_PRO])
+
+        # Local backends always available (if dependencies installed)
+        try:
+            from diffusers import Flux2KleinPipeline
+            available.append(BackendType.FLUX_KLEIN)
+        except ImportError:
+            pass
+
+        try:
+            from diffusers import ZImagePipeline
+            available.append(BackendType.ZIMAGE_TURBO)
+            available.append(BackendType.ZIMAGE_BASE)
+        except ImportError:
+            pass
+
+        try:
+            from diffusers import LongCatImageEditPipeline
+            available.append(BackendType.LONGCAT_EDIT)
+        except ImportError:
+            pass
+
+        try:
+            from diffusers import QwenImageEditPlusPipeline
+            available.append(BackendType.QWEN_IMAGE_EDIT)
+        except ImportError:
+            pass
+
+        # ComfyUI backend - check if ComfyUI client works
+        try:
+            from .comfyui_client import ComfyUIClient
+            client = ComfyUIClient()
+            if client.is_healthy():
+                available.append(BackendType.QWEN_COMFYUI)
+        except Exception:
+            pass
+
+        return available
+
+    @staticmethod
+    def get_backend_choices() -> list:
+        """Get list of backend choices for UI dropdowns."""
+        return [
+            ("Gemini Flash (Cloud)", BackendType.GEMINI_FLASH.value),
+            ("Gemini Pro (Cloud)", BackendType.GEMINI_PRO.value),
+            ("FLUX.2 klein 4B (Local)", BackendType.FLUX_KLEIN.value),
+            ("Z-Image Turbo 6B (Fast, 9 steps, 16GB)", BackendType.ZIMAGE_TURBO.value),
+            ("Z-Image Base 6B (Quality, 50 steps, CFG)", BackendType.ZIMAGE_BASE.value),
+            ("LongCat-Image-Edit (Instruction Editing, 18GB)", BackendType.LONGCAT_EDIT.value),
+            ("Qwen-Image-Edit-2511 (Local, High VRAM)", BackendType.QWEN_IMAGE_EDIT.value),
+            ("Qwen-Image-Edit-2511-FP8 (ComfyUI)", BackendType.QWEN_COMFYUI.value),
+        ]
+
+    @staticmethod
+    def backend_from_string(value: str) -> BackendType:
+        """Convert string to BackendType."""
+        for bt in BackendType:
+            if bt.value == value:
+                return bt
+        raise ValueError(f"Unknown backend: {value}")

src/character_service.py

@@ -0,0 +1,709 @@
+"""
+Character Sheet Service
+=======================
+
+9-stage pipeline for generating 7-view character turnaround sheets.
+
+Layout:
++------------------+------------------+------------------+
+| Left Face Profile| Front Face       | Right Face Profile|  (3:4)
++------------------+------------------+------------------+
+| Left Side Body   | Front Body | Right Side Body | Back Body |  (9:16)
++------------------+------------------+------------------+
+"""
+
+import time
+import random
+import logging
+from pathlib import Path
+from typing import Optional, Tuple, Dict, Any, Callable, List
+from datetime import datetime
+from PIL import Image
+
+from .models import (
+    GenerationRequest,
+    GenerationResult,
+    CharacterSheetConfig,
+    CharacterSheetMetadata
+)
+from .gemini_client import GeminiClient
+from .backend_router import BackendRouter, BackendType
+from .utils import ensure_pil_image, save_image, sanitize_filename, preprocess_images_for_backend
+
+
+logger = logging.getLogger(__name__)
+
+
+class CharacterSheetService:
+    """
+    Service for generating 7-view character turnaround sheets.
+
+    Pipeline (9 stages):
+    0. Input normalization (face→body or body→face+body)
+    1. Front face portrait
+    2. Left face profile (90 degrees)
+    3. Right face profile (90 degrees)
+    4. Front full body (from normalized)
+    5. Back full body
+    6. Left side full body
+    7. Right side full body
+    8. Composite character sheet
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        use_pro_model: bool = False,
+        config: Optional[CharacterSheetConfig] = None,
+        backend: Optional[BackendType] = None,
+        backend_router: Optional[BackendRouter] = None
+    ):
+        """
+        Initialize character sheet service.
+
+        Args:
+            api_key: Gemini API key (for cloud backends)
+            use_pro_model: Use Gemini Pro model (legacy, use backend param instead)
+            config: Optional configuration
+            backend: Specific backend to use
+            backend_router: Pre-configured backend router
+        """
+        self.config = config or CharacterSheetConfig()
+
+        # Determine backend
+        if backend_router is not None:
+            self.router = backend_router
+            self.backend = backend or backend_router.default_backend
+        else:
+            # Determine default backend based on params
+            if backend is not None:
+                self.backend = backend
+            elif use_pro_model:
+                self.backend = BackendType.GEMINI_PRO
+            else:
+                self.backend = BackendType.GEMINI_FLASH
+
+            self.router = BackendRouter(
+                gemini_api_key=api_key,
+                default_backend=self.backend
+            )
+
+        # For backward compatibility
+        self.use_pro_model = use_pro_model
+        self.client = self.router.get_client(self.backend)
+
+        logger.info(f"CharacterSheetService initialized (backend: {self.backend.value})")
+
+    def generate_character_sheet(
+        self,
+        initial_image: Optional[Image.Image],
+        input_type: str = "Face Only",
+        character_name: str = "Character",
+        gender_term: str = "character",
+        costume_description: str = "",
+        costume_image: Optional[Image.Image] = None,
+        face_image: Optional[Image.Image] = None,
+        body_image: Optional[Image.Image] = None,
+        include_costume_in_faces: bool = True,
+        progress_callback: Optional[Callable[[int, int, str], None]] = None,
+        stage_callback: Optional[Callable[[str, Image.Image, Dict[str, Any]], None]] = None,
+        output_dir: Optional[Path] = None
+    ) -> Tuple[Optional[Image.Image], str, Dict[str, Any]]:
+        """
+        Generate complete 7-view character turnaround sheet.
+
+        Args:
+            initial_image: Starting image (face or body)
+            input_type: "Face Only", "Full Body", or "Face + Body (Separate)"
+            character_name: Character name
+            gender_term: "character", "man", or "woman"
+            costume_description: Text costume description
+            costume_image: Optional costume reference
+            face_image: Face image (for Face + Body mode)
+            body_image: Body image (for Face + Body mode)
+            include_costume_in_faces: If True, include costume reference in face views.
+                Set False for models like FLUX that confuse costume with framing.
+            progress_callback: Optional callback(stage, total_stages, message)
+            stage_callback: Optional callback(stage_name, image, stages_dict) called after each
+                stage completes with the generated image. Enables streaming preview.
+            output_dir: Optional output directory
+
+        Returns:
+            Tuple of (character_sheet, status_message, metadata)
+        """
+        try:
+            total_stages = 9
+            stages = {}
+
+            logger.info("=" * 60)
+            logger.info(f"STARTING CHARACTER SHEET: {character_name}")
+            logger.info(f"Input type: {input_type}")
+            logger.info(f"Costume: {costume_description or '(none)'}")
+            logger.info("=" * 60)
+
+            # Build costume instructions - separate for face and body views
+            # For models like FLUX, costume refs confuse face generation
+            costume_instruction_body = ""
+            if costume_description:
+                costume_instruction_body = f" wearing {costume_description}"
+            elif costume_image:
+                costume_instruction_body = " wearing the costume shown in the reference"
+
+            # Face views only get costume instruction if flag is set
+            if include_costume_in_faces:
+                costume_instruction_face = costume_instruction_body
+            else:
+                costume_instruction_face = ""
+                logger.info("Costume excluded from face views (include_costume_in_faces=False)")
+
+            def update_progress(stage: int, message: str):
+                if progress_callback:
+                    progress_callback(stage, total_stages, message)
+                logger.info(f"[Stage {stage}/{total_stages}] {message}")
+
+            def notify_stage_complete(stage_name: str, image: Image.Image):
+                """Notify callback when a stage completes for streaming preview."""
+                if stage_callback and image is not None:
+                    stage_callback(stage_name, image, stages)
+
+            # =================================================================
+            # Stage 0: Normalize input
+            # =================================================================
+            update_progress(0, "Normalizing input images...")
+
+            reference_body, reference_face = self._normalize_input(
+                initial_image=initial_image,
+                input_type=input_type,
+                face_image=face_image,
+                body_image=body_image,
+                costume_instruction=costume_instruction_body,  # Body normalization uses full costume
+                costume_image=costume_image,
+                gender_term=gender_term,
+                stages=stages,
+                progress_callback=lambda msg: update_progress(0, msg)
+            )
+
+            if reference_body is None or reference_face is None:
+                return None, "Failed to normalize input images", {}
+
+            time.sleep(1)
+
+            # =================================================================
+            # FACE VIEWS (3 portraits)
+            # =================================================================
+
+            # Stage 1: Front face portrait
+            update_progress(1, "Generating front face portrait...")
+
+            if input_type == "Face + Body (Separate)":
+                prompt = f"Generate a close-up frontal facial portrait showing the {gender_term} from the first image (body/costume reference), extrapolate and extract exact facial details from the second image (face reference). Do NOT transfer clothing or hair style from the second image. The face should fill the entire vertical space, neutral grey background with professional studio lighting."
+                input_images = [reference_body, reference_face]
+            else:
+                prompt = f"Generate a formal portrait view of this {gender_term}{costume_instruction_face} as depicted in the reference images, in front of a neutral grey background with professional studio lighting. The face should fill the entire vertical space. Maintain exact facial features from the reference."
+                input_images = [reference_face, reference_body]
+                # Only include costume in face views if flag is set (smarter models)
+                if costume_image and include_costume_in_faces:
+                    input_images.append(costume_image)
+
+            front_face, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.face_aspect_ratio,
+                temperature=self.config.face_temperature
+            )
+
+            if front_face is None:
+                return None, f"Stage 1 failed: {status}", {}
+
+            stages['front_face'] = front_face
+            notify_stage_complete('front_face', front_face)
+            time.sleep(1)
+
+            # Stage 2: Left face profile
+            update_progress(2, "Generating left face profile...")
+
+            prompt = f"Create a left side profile view (90 degrees) of this {gender_term}'s face{costume_instruction_face}, showing the left side of the face filling the frame. Professional studio lighting against a neutral grey background. Maintain exact facial features from the reference."
+
+            input_images = [front_face, reference_body]
+            if input_type == "Face + Body (Separate)":
+                input_images.append(reference_face)
+            elif costume_image and include_costume_in_faces:
+                # Only include costume in face views if flag is set (smarter models)
+                input_images.append(costume_image)
+
+            left_face, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.face_aspect_ratio,
+                temperature=self.config.face_temperature
+            )
+
+            if left_face is None:
+                return None, f"Stage 2 failed: {status}", {}
+
+            stages['left_face'] = left_face
+            notify_stage_complete('left_face', left_face)
+            time.sleep(1)
+
+            # Stage 3: Right face profile
+            update_progress(3, "Generating right face profile...")
+
+            prompt = f"Create a right side profile view (90 degrees) of this {gender_term}'s face{costume_instruction_face}, showing the right side of the face filling the frame. Professional studio lighting against a neutral grey background. Maintain exact facial features from the reference."
+
+            input_images = [front_face, reference_body]
+            if input_type == "Face + Body (Separate)":
+                input_images.append(reference_face)
+            elif costume_image and include_costume_in_faces:
+                # Only include costume in face views if flag is set (smarter models)
+                input_images.append(costume_image)
+
+            right_face, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.face_aspect_ratio,
+                temperature=self.config.face_temperature
+            )
+
+            if right_face is None:
+                return None, f"Stage 3 failed: {status}", {}
+
+            stages['right_face'] = right_face
+            notify_stage_complete('right_face', right_face)
+            time.sleep(1)
+
+            # =================================================================
+            # BODY VIEWS (4 views)
+            # =================================================================
+
+            # Stage 4: Front body (use normalized reference)
+            update_progress(4, "Using front body from normalized reference...")
+            front_body = reference_body
+            stages['front_body'] = front_body
+            notify_stage_complete('front_body', front_body)
+            time.sleep(1)
+
+            # Stage 5: Back body
+            update_progress(5, "Generating back full body...")
+
+            prompt = f"Generate a rear view image of this {gender_term}{costume_instruction_body} showing the back in a neutral standing pose against a neutral grey background with professional studio lighting. The full body should fill the vertical space. Maintain consistent appearance from the reference images."
+
+            input_images = [reference_body, front_face]
+            if costume_image:
+                input_images.append(costume_image)
+
+            back_body, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.body_temperature
+            )
+
+            if back_body is None:
+                return None, f"Stage 5 failed: {status}", {}
+
+            stages['back_body'] = back_body
+            notify_stage_complete('back_body', back_body)
+            time.sleep(1)
+
+            # Stage 6: Left side body
+            update_progress(6, "Generating left side full body...")
+
+            prompt = f"Generate a left side view of the full body of this {gender_term}{costume_instruction_body} in front of a neutral grey background. The {gender_term} should be shown from the left side (90 degree angle) in a neutral standing pose. Full body fills vertical space. Professional studio lighting."
+
+            input_images = [left_face, front_body, reference_body]
+
+            left_body, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.body_temperature
+            )
+
+            if left_body is None:
+                return None, f"Stage 6 failed: {status}", {}
+
+            stages['left_body'] = left_body
+            notify_stage_complete('left_body', left_body)
+            time.sleep(1)
+
+            # Stage 7: Right side body
+            update_progress(7, "Generating right side full body...")
+
+            prompt = f"Generate a right side view of the full body of this {gender_term}{costume_instruction_body} in front of a neutral grey background. The {gender_term} should be shown from the right side (90 degree angle) in a neutral standing pose. Full body fills vertical space. Professional studio lighting."
+
+            input_images = [right_face, front_body, reference_body]
+
+            right_body, status = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.body_temperature
+            )
+
+            if right_body is None:
+                return None, f"Stage 7 failed: {status}", {}
+
+            stages['right_body'] = right_body
+            notify_stage_complete('right_body', right_body)
+            time.sleep(1)
+
+            # =================================================================
+            # Stage 8: Composite character sheet
+            # =================================================================
+            update_progress(8, "Compositing character sheet...")
+
+            character_sheet = self.composite_character_sheet(
+                left_face=left_face,
+                front_face=front_face,
+                right_face=right_face,
+                left_body=left_body,
+                front_body=front_body,
+                right_body=right_body,
+                back_body=back_body
+            )
+
+            stages['character_sheet'] = character_sheet
+
+            # Build metadata
+            metadata = CharacterSheetMetadata(
+                character_name=character_name,
+                input_type=input_type,
+                costume_description=costume_description,
+                backend=self.router.get_active_backend_name(),
+                stages={
+                    "left_face": {"size": left_face.size},
+                    "front_face": {"size": front_face.size},
+                    "right_face": {"size": right_face.size},
+                    "left_body": {"size": left_body.size},
+                    "front_body": {"size": front_body.size},
+                    "right_body": {"size": right_body.size},
+                    "back_body": {"size": back_body.size},
+                }
+            )
+
+            success_msg = f"Character sheet generated! 7 views of {character_name}"
+
+            # Save to disk if requested
+            if output_dir:
+                save_dir = self._save_outputs(
+                    character_name=character_name,
+                    stages=stages,
+                    output_dir=output_dir
+                )
+                success_msg += f"\nSaved to: {save_dir}"
+
+            update_progress(9, "Complete!")
+            return character_sheet, success_msg, {"metadata": metadata, "stages": stages}
+
+        except Exception as e:
+            logger.exception(f"Character sheet generation failed: {e}")
+            return None, f"Error: {str(e)}", {}
+
+    def _normalize_input(
+        self,
+        initial_image: Optional[Image.Image],
+        input_type: str,
+        face_image: Optional[Image.Image],
+        body_image: Optional[Image.Image],
+        costume_instruction: str,
+        costume_image: Optional[Image.Image],
+        gender_term: str,
+        stages: dict,
+        progress_callback: Optional[Callable]
+    ) -> Tuple[Optional[Image.Image], Optional[Image.Image]]:
+        """Normalize input images to create reference body and face."""
+
+        if input_type == "Face + Body (Separate)":
+            if face_image is None or body_image is None:
+                return None, None
+
+            if progress_callback:
+                progress_callback("Normalizing body image...")
+
+            prompt = f"Front view full body portrait of this person{costume_instruction}, standing, neutral background"
+            input_images = [body_image, face_image]
+            if costume_image:
+                input_images.append(costume_image)
+
+            normalized_body, _ = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.normalize_temperature
+            )
+
+            if normalized_body is None:
+                return None, None
+
+            stages['normalized_body'] = normalized_body
+            return normalized_body, face_image
+
+        elif input_type == "Face Only":
+            if initial_image is None:
+                return None, None
+
+            if progress_callback:
+                progress_callback("Generating full body from face...")
+
+            prompt = f"Create a full body image of the {gender_term}{costume_instruction} standing in a neutral pose in front of a grey background with professional studio lighting. The {gender_term}'s face and features should match the reference image exactly."
+
+            input_images = [initial_image]
+            if costume_image:
+                input_images.append(costume_image)
+
+            full_body, _ = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.normalize_temperature
+            )
+
+            if full_body is None:
+                return None, None
+
+            stages['generated_body'] = full_body
+            return full_body, initial_image
+
+        else:  # Full Body
+            if initial_image is None:
+                return None, None
+
+            # Normalize body
+            if progress_callback:
+                progress_callback("Normalizing full body...")
+
+            prompt = f"Front view full body portrait of this person{costume_instruction}, standing, neutral background"
+
+            input_images = [initial_image]
+            if costume_image:
+                input_images.append(costume_image)
+
+            normalized_body, _ = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.body_aspect_ratio,
+                temperature=self.config.normalize_temperature
+            )
+
+            if normalized_body is None:
+                return None, None
+
+            stages['normalized_body'] = normalized_body
+            time.sleep(1)
+
+            # Extract face
+            if progress_callback:
+                progress_callback("Generating face closeup...")
+
+            prompt = f"Create a frontal closeup portrait of this {gender_term}'s face{costume_instruction}, focusing only on the face and head. Professional studio lighting against a neutral grey background. The face should fill the entire vertical space. Maintain exact facial features from the reference."
+
+            input_images = [normalized_body, initial_image]
+            if costume_image:
+                input_images.append(costume_image)
+
+            face_closeup, _ = self._generate_stage(
+                prompt=prompt,
+                input_images=input_images,
+                aspect_ratio=self.config.face_aspect_ratio,
+                temperature=self.config.face_temperature
+            )
+
+            if face_closeup is None:
+                return None, None
+
+            stages['extracted_face'] = face_closeup
+            return normalized_body, face_closeup
+
+    def _generate_stage(
+        self,
+        prompt: str,
+        input_images: List[Image.Image],
+        aspect_ratio: str,
+        temperature: float,
+        max_retries: int = 3
+    ) -> Tuple[Optional[Image.Image], str]:
+        """Generate single stage with retry logic."""
+
+        modified_prompt = prompt
+        cfg = self.config
+
+        # Preprocess images for the current backend
+        backend_type = self.backend.value if self.backend else "unknown"
+        processed_images = preprocess_images_for_backend(
+            input_images, backend_type, aspect_ratio
+        )
+        logger.info(f"Preprocessed {len(processed_images)} images for {backend_type}")
+
+        for attempt in range(max_retries):
+            try:
+                if attempt > 0:
+                    wait_time = cfg.retry_delay
+                    logger.info(f"Retry {attempt + 1}/{max_retries}, waiting {wait_time}s...")
+                    time.sleep(wait_time)
+
+                request = GenerationRequest(
+                    prompt=modified_prompt,
+                    input_images=processed_images,
+                    aspect_ratio=aspect_ratio,
+                    temperature=temperature
+                )
+
+                result = self.client.generate(request)
+
+                if result.success:
+                    delay = random.uniform(cfg.rate_limit_delay_min, cfg.rate_limit_delay_max)
+                    time.sleep(delay)
+                    return result.image, result.message
+
+                # Check for safety block
+                error_upper = result.message.upper()
+                if any(kw in error_upper for kw in ['SAFETY', 'BLOCKED', 'PROHIBITED', 'IMAGE_OTHER']):
+                    if 'wearing' not in modified_prompt.lower():
+                        if 'body' in modified_prompt.lower():
+                            modified_prompt = prompt + ", fully clothed in casual wear"
+                        else:
+                            modified_prompt = prompt + ", wearing appropriate clothing"
+                        logger.info("Modified prompt to avoid safety filters")
+
+                logger.warning(f"Attempt {attempt + 1} failed: {result.message}")
+
+            except Exception as e:
+                logger.error(f"Attempt {attempt + 1} exception: {e}")
+                if attempt == max_retries - 1:
+                    return None, str(e)
+
+        return None, f"All {max_retries} attempts failed"
+
+    def composite_character_sheet(
+        self,
+        left_face: Image.Image,
+        front_face: Image.Image,
+        right_face: Image.Image,
+        left_body: Image.Image,
+        front_body: Image.Image,
+        right_body: Image.Image,
+        back_body: Image.Image
+    ) -> Image.Image:
+        """
+        Composite all 7 views into character sheet.
+
+        Layout:
+        +------------------+------------------+------------------+
+        | Left Face Profile| Front Face       | Right Face Profile|
+        +------------------+------------------+------------------+
+        | Left Side Body   | Front Body | Right Side Body | Back Body |
+        +------------------+------------------+------------------+
+        """
+        # Normalize all inputs
+        left_face = ensure_pil_image(left_face, "left_face")
+        front_face = ensure_pil_image(front_face, "front_face")
+        right_face = ensure_pil_image(right_face, "right_face")
+        left_body = ensure_pil_image(left_body, "left_body")
+        front_body = ensure_pil_image(front_body, "front_body")
+        right_body = ensure_pil_image(right_body, "right_body")
+        back_body = ensure_pil_image(back_body, "back_body")
+
+        spacing = self.config.spacing
+
+        # Calculate dimensions
+        face_row_width = left_face.width + front_face.width + right_face.width
+        body_row_width = left_body.width + front_body.width + right_body.width + back_body.width
+        canvas_width = max(face_row_width, body_row_width)
+        canvas_height = front_face.height + spacing + front_body.height
+
+        # Create canvas
+        canvas = Image.new('RGB', (canvas_width, canvas_height), color=self.config.background_color)
+
+        # Upper row: 3 face portraits
+        x = 0
+        canvas.paste(left_face, (x, 0))
+        x += left_face.width
+        canvas.paste(front_face, (x, 0))
+        x += front_face.width
+        canvas.paste(right_face, (x, 0))
+
+        # Lower row: 4 body views
+        x = 0
+        y = front_face.height + spacing
+        canvas.paste(left_body, (x, y))
+        x += left_body.width
+        canvas.paste(front_body, (x, y))
+        x += front_body.width
+        canvas.paste(right_body, (x, y))
+        x += right_body.width
+        canvas.paste(back_body, (x, y))
+
+        return canvas
+
+    def extract_views_from_sheet(
+        self,
+        character_sheet: Image.Image
+    ) -> Dict[str, Image.Image]:
+        """
+        Extract individual views from character sheet.
+
+        Returns:
+            Dictionary with 7 extracted views
+        """
+        sheet_width, sheet_height = character_sheet.size
+        spacing = self.config.spacing
+
+        # Find separator by scanning for dark bar
+        scan_start = sheet_height // 3
+        scan_end = (2 * sheet_height) // 3
+
+        min_brightness = 255
+        separator_y = scan_start
+
+        for y in range(scan_start, scan_end):
+            line = character_sheet.crop((0, y, min(200, sheet_width), y + 1))
+            pixels = list(line.getdata())
+            avg_brightness = sum(
+                sum(p[:3]) / 3 if isinstance(p, tuple) else p
+                for p in pixels
+            ) / len(pixels)
+
+            if avg_brightness < min_brightness:
+                min_brightness = avg_brightness
+                separator_y = y
+
+        face_height = separator_y
+        body_start_y = separator_y + spacing
+        body_height = sheet_height - body_start_y
+
+        # Calculate widths from aspect ratios
+        face_width = (face_height * 3) // 4
+        body_width = (body_height * 9) // 16
+
+        # Extract views
+        views = {
+            'left_face': character_sheet.crop((0, 0, face_width, face_height)),
+            'front_face': character_sheet.crop((face_width, 0, 2 * face_width, face_height)),
+            'right_face': character_sheet.crop((2 * face_width, 0, 3 * face_width, face_height)),
+            'left_body': character_sheet.crop((0, body_start_y, body_width, body_start_y + body_height)),
+            'front_body': character_sheet.crop((body_width, body_start_y, 2 * body_width, body_start_y + body_height)),
+            'right_body': character_sheet.crop((2 * body_width, body_start_y, 3 * body_width, body_start_y + body_height)),
+            'back_body': character_sheet.crop((3 * body_width, body_start_y, 4 * body_width, body_start_y + body_height)),
+        }
+
+        return views
+
+    def _save_outputs(
+        self,
+        character_name: str,
+        stages: dict,
+        output_dir: Path
+    ) -> Path:
+        """Save all outputs to directory."""
+        output_dir = Path(output_dir)
+        safe_name = sanitize_filename(character_name)
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        char_dir = output_dir / f"{safe_name}_{timestamp}"
+        char_dir.mkdir(parents=True, exist_ok=True)
+
+        for name, image in stages.items():
+            if isinstance(image, Image.Image):
+                save_image(image, char_dir, f"{safe_name}_{name}")
+
+        logger.info(f"Saved outputs to: {char_dir}")
+        return char_dir

src/comfyui_client.py

@@ -0,0 +1,578 @@
+"""
+ComfyUI Client for Qwen-Image-Edit-2511
+========================================
+
+Client to interact with ComfyUI API for running Qwen-Image-Edit-2511.
+
+Model setup (download from HuggingFace):
+
+  Lightning (default, 4-step):
+    diffusion_models/  qwen_image_edit_2511_fp8_e4m3fn_scaled_lightning_comfyui_4steps_v1.0.safetensors
+                         (lightx2v/Qwen-Image-Edit-2511-Lightning)
+
+  Standard (20-step, optional):
+    diffusion_models/  qwen_image_edit_2511_fp8mixed.safetensors
+                         (Comfy-Org/Qwen-Image-Edit_ComfyUI)
+
+  Shared:
+    text_encoders/     qwen_2.5_vl_7b_fp8_scaled.safetensors         (Comfy-Org/Qwen-Image_ComfyUI)
+    vae/               qwen_image_vae.safetensors                     (Comfy-Org/Qwen-Image_ComfyUI)
+
+Required custom nodes:
+  - Comfyui-QwenEditUtils (lrzjason) for TextEncodeQwenImageEditPlus
+"""
+
+import logging
+import time
+import uuid
+import json
+import io
+import base64
+from typing import Optional, List, Tuple
+from PIL import Image
+import websocket
+import urllib.request
+import urllib.parse
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class ComfyUIClient:
+    """
+    Client for ComfyUI API to run Qwen-Image-Edit-2511.
+
+    Requires ComfyUI running with:
+    - Qwen-Image-Edit-2511 model in models/diffusion_models/
+    - Qwen 2.5 VL 7B text encoder in models/text_encoders/
+    - Qwen Image VAE in models/vae/
+    - Comfyui-QwenEditUtils custom node installed
+    """
+
+    # Default ComfyUI settings
+    DEFAULT_HOST = "127.0.0.1"
+    DEFAULT_PORT = 8188
+
+    # Model file names (expected in ComfyUI models/ subfolders)
+    # Lightning: baked model (LoRA pre-merged, ComfyUI-specific format)
+    UNET_MODEL_LIGHTNING = "qwen_image_edit_2511_fp8_e4m3fn_scaled_lightning_comfyui_4steps_v1.0.safetensors"
+    # Standard: base fp8mixed model (20-step, higher quality)
+    UNET_MODEL_STANDARD = "qwen_image_edit_2511_fp8mixed.safetensors"
+    TEXT_ENCODER = "qwen_2.5_vl_7b_fp8_scaled.safetensors"
+    VAE_MODEL = "qwen_image_vae.safetensors"
+
+    # Target output dimensions per aspect ratio.
+    # Generation happens at 1024x1024, then crop+resize to these.
+    ASPECT_RATIOS = {
+        "1:1": (1024, 1024),
+        "16:9": (1344, 768),
+        "9:16": (768, 1344),
+        "21:9": (1680, 720),
+        "3:2": (1248, 832),
+        "2:3": (832, 1248),
+        "3:4": (896, 1152),
+        "4:3": (1152, 896),
+        "4:5": (1024, 1280),
+        "5:4": (1280, 1024),
+    }
+
+    # Generate at 1024x1024 (proven safe for Qwen's VAE), then crop+resize
+    NATIVE_RESOLUTION = (1024, 1024)
+
+    # With Lightning LoRA: 4 steps, CFG 1.0  (fast, ~seconds per view)
+    # Without LoRA:        20 steps, CFG 4.0
+    DEFAULT_STEPS_LIGHTNING = 4
+    DEFAULT_STEPS_STANDARD = 20
+    DEFAULT_CFG_LIGHTNING = 1.0
+    DEFAULT_CFG_STANDARD = 4.0
+
+    def __init__(
+        self,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+        use_lightning: bool = True,
+    ):
+        """
+        Initialize ComfyUI client.
+
+        Args:
+            host: ComfyUI server host
+            port: ComfyUI server port
+            use_lightning: Use Lightning LoRA for 4-step generation (much faster)
+        """
+        self.host = host
+        self.port = port
+        self.use_lightning = use_lightning
+        self.client_id = str(uuid.uuid4())
+        self.server_address = f"{host}:{port}"
+
+        if use_lightning:
+            self.num_inference_steps = self.DEFAULT_STEPS_LIGHTNING
+            self.cfg_scale = self.DEFAULT_CFG_LIGHTNING
+        else:
+            self.num_inference_steps = self.DEFAULT_STEPS_STANDARD
+            self.cfg_scale = self.DEFAULT_CFG_STANDARD
+
+        logger.info(
+            f"ComfyUIClient initialized for {self.server_address} "
+            f"(lightning={use_lightning}, steps={self.num_inference_steps})"
+        )
+
+    def is_healthy(self) -> bool:
+        """Check if ComfyUI server is running and accessible."""
+        try:
+            url = f"http://{self.server_address}/system_stats"
+            with urllib.request.urlopen(url, timeout=5) as response:
+                return response.status == 200
+        except Exception:
+            return False
+
+    def _upload_image(self, image: Image.Image, name: str = "input.png") -> Optional[str]:
+        """
+        Upload an image to ComfyUI, pre-resized to fit within 1024x1024.
+
+        Args:
+            image: PIL Image to upload
+            name: Filename for the uploaded image
+
+        Returns:
+            Filename on server, or None if failed
+        """
+        try:
+            # Pre-resize to keep total pixels around 1024x1024 (matching reference workflow)
+            max_pixels = 1024 * 1024
+            w, h = image.size
+            if w * h > max_pixels:
+                scale = (max_pixels / (w * h)) ** 0.5
+                new_w = int(w * scale)
+                new_h = int(h * scale)
+                image = image.resize((new_w, new_h), Image.LANCZOS)
+                logger.debug(f"Pre-resized input from {w}x{h} to {new_w}x{new_h}")
+
+            # Convert image to bytes
+            img_bytes = io.BytesIO()
+            image.save(img_bytes, format='PNG')
+            img_bytes.seek(0)
+
+            # Create multipart form data
+            boundary = uuid.uuid4().hex
+
+            body = b''
+            body += f'--{boundary}\r\n'.encode()
+            body += f'Content-Disposition: form-data; name="image"; filename="{name}"\r\n'.encode()
+            body += b'Content-Type: image/png\r\n\r\n'
+            body += img_bytes.read()
+            body += f'\r\n--{boundary}--\r\n'.encode()
+
+            url = f"http://{self.server_address}/upload/image"
+            req = urllib.request.Request(
+                url,
+                data=body,
+                headers={
+                    'Content-Type': f'multipart/form-data; boundary={boundary}'
+                }
+            )
+
+            with urllib.request.urlopen(req) as response:
+                result = json.loads(response.read())
+                return result.get('name')
+
+        except Exception as e:
+            logger.error(f"Failed to upload image: {e}")
+            return None
+
+    def _queue_prompt(self, prompt: dict) -> str:
+        """
+        Queue a prompt for execution.
+
+        Args:
+            prompt: Workflow prompt dict
+
+        Returns:
+            Prompt ID
+        """
+        prompt_id = str(uuid.uuid4())
+        p = {"prompt": prompt, "client_id": self.client_id, "prompt_id": prompt_id}
+        data = json.dumps(p).encode('utf-8')
+
+        url = f"http://{self.server_address}/prompt"
+        req = urllib.request.Request(url, data=data)
+        urllib.request.urlopen(req)
+
+        return prompt_id
+
+    def _get_history(self, prompt_id: str) -> dict:
+        """Get execution history for a prompt."""
+        url = f"http://{self.server_address}/history/{prompt_id}"
+        with urllib.request.urlopen(url) as response:
+            return json.loads(response.read())
+
+    def _get_image(self, filename: str, subfolder: str, folder_type: str) -> bytes:
+        """Get an image from ComfyUI."""
+        data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+        url_values = urllib.parse.urlencode(data)
+        url = f"http://{self.server_address}/view?{url_values}"
+        with urllib.request.urlopen(url) as response:
+            return response.read()
+
+    def _wait_for_completion(self, prompt_id: str, timeout: float = 900.0) -> bool:
+        """
+        Wait for prompt execution to complete using websocket.
+
+        Args:
+            prompt_id: The prompt ID to wait for
+            timeout: Maximum time to wait in seconds (default 15 min for image editing)
+
+        Returns:
+            True if completed successfully, False if timeout/error
+        """
+        ws = None
+        try:
+            ws_url = f"ws://{self.server_address}/ws?clientId={self.client_id}"
+            ws = websocket.WebSocket()
+            ws.settimeout(timeout)
+            ws.connect(ws_url)
+
+            start_time = time.time()
+            while time.time() - start_time < timeout:
+                try:
+                    out = ws.recv()
+                    if isinstance(out, str):
+                        message = json.loads(out)
+                        if message['type'] == 'executing':
+                            data = message['data']
+                            if data['node'] is None and data['prompt_id'] == prompt_id:
+                                return True  # Execution complete
+                        elif message['type'] == 'execution_error':
+                            logger.error(f"Execution error: {message}")
+                            return False
+                except websocket.WebSocketTimeoutException:
+                    continue
+
+            logger.error("Timeout waiting for completion")
+            return False
+
+        except Exception as e:
+            logger.error(f"WebSocket error: {e}")
+            return False
+        finally:
+            if ws:
+                try:
+                    ws.close()
+                except:
+                    pass
+
+    def _get_dimensions(self, aspect_ratio: str) -> Tuple[int, int]:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return self.ASPECT_RATIOS.get(ratio, (1024, 1024))
+
+    @staticmethod
+    def _crop_and_resize(image: Image.Image, target_w: int, target_h: int) -> Image.Image:
+        """Crop to target aspect ratio, then resize. Centers the crop."""
+        src_w, src_h = image.size
+        target_ratio = target_w / target_h
+        src_ratio = src_w / src_h
+
+        if abs(target_ratio - src_ratio) < 0.01:
+            return image.resize((target_w, target_h), Image.LANCZOS)
+
+        if target_ratio < src_ratio:
+            crop_w = int(src_h * target_ratio)
+            offset = (src_w - crop_w) // 2
+            image = image.crop((offset, 0, offset + crop_w, src_h))
+        else:
+            crop_h = int(src_w / target_ratio)
+            offset = (src_h - crop_h) // 2
+            image = image.crop((0, offset, src_w, offset + crop_h))
+
+        return image.resize((target_w, target_h), Image.LANCZOS)
+
+    def _build_workflow(
+        self,
+        prompt: str,
+        width: int,
+        height: int,
+        input_images: List[str] = None,
+        negative_prompt: str = ""
+    ) -> dict:
+        """
+        Build the ComfyUI workflow for Qwen-Image-Edit-2511.
+
+        Workflow graph:
+          UNETLoader → KSampler
+          CLIPLoader → TextEncodeQwenImageEditPlus (pos/neg)
+          VAELoader → TextEncode + VAEDecode
+          LoadImage(s) → TextEncodeQwenImageEditPlus
+          EmptyQwenImageLayeredLatentImage → KSampler
+          KSampler → VAEDecode → PreviewImage
+
+        Lightning mode uses a baked model (LoRA pre-merged), no separate
+        LoRA or ModelSamplingAuraFlow nodes needed.
+        """
+        workflow = {}
+        node_id = 1
+
+        # --- Model loading ---
+
+        # Select model based on lightning mode
+        unet_name = (self.UNET_MODEL_LIGHTNING if self.use_lightning
+                     else self.UNET_MODEL_STANDARD)
+
+        # UNETLoader - weight_dtype "default" lets ComfyUI auto-detect fp8
+        unet_id = str(node_id)
+        workflow[unet_id] = {
+            "class_type": "UNETLoader",
+            "inputs": {
+                "unet_name": unet_name,
+                "weight_dtype": "default"
+            }
+        }
+        node_id += 1
+
+        # CLIPLoader
+        clip_id = str(node_id)
+        workflow[clip_id] = {
+            "class_type": "CLIPLoader",
+            "inputs": {
+                "clip_name": self.TEXT_ENCODER,
+                "type": "qwen_image"
+            }
+        }
+        node_id += 1
+
+        # VAELoader
+        vae_id = str(node_id)
+        workflow[vae_id] = {
+            "class_type": "VAELoader",
+            "inputs": {
+                "vae_name": self.VAE_MODEL
+            }
+        }
+        node_id += 1
+
+        model_out_id = unet_id
+
+        # --- Input images ---
+
+        image_loader_ids = []
+        if input_images:
+            for img_name in input_images[:3]:  # Max 3 reference images
+                img_loader_id = str(node_id)
+                workflow[img_loader_id] = {
+                    "class_type": "LoadImage",
+                    "inputs": {
+                        "image": img_name
+                    }
+                }
+                image_loader_ids.append(img_loader_id)
+                node_id += 1
+
+        # --- Text encoding ---
+
+        # Positive: prompt + vision references + VAE
+        pos_encode_id = str(node_id)
+        pos_inputs = {
+            "clip": [clip_id, 0],
+            "prompt": prompt,
+            "vae": [vae_id, 0]
+        }
+        for i, loader_id in enumerate(image_loader_ids):
+            pos_inputs[f"image{i+1}"] = [loader_id, 0]
+
+        workflow[pos_encode_id] = {
+            "class_type": "TextEncodeQwenImageEditPlus",
+            "inputs": pos_inputs
+        }
+        node_id += 1
+
+        # Negative: text only, no images
+        neg_encode_id = str(node_id)
+        workflow[neg_encode_id] = {
+            "class_type": "TextEncodeQwenImageEditPlus",
+            "inputs": {
+                "clip": [clip_id, 0],
+                "prompt": negative_prompt or " ",
+                "vae": [vae_id, 0]
+            }
+        }
+        node_id += 1
+
+        # --- Latent + sampling ---
+
+        latent_id = str(node_id)
+        workflow[latent_id] = {
+            "class_type": "EmptySD3LatentImage",
+            "inputs": {
+                "width": width,
+                "height": height,
+                "batch_size": 1
+            }
+        }
+        node_id += 1
+
+        sampler_id = str(node_id)
+        workflow[sampler_id] = {
+            "class_type": "KSampler",
+            "inputs": {
+                "model": [model_out_id, 0],
+                "positive": [pos_encode_id, 0],
+                "negative": [neg_encode_id, 0],
+                "latent_image": [latent_id, 0],
+                "seed": int(time.time()) % 2**32,
+                "steps": self.num_inference_steps,
+                "cfg": self.cfg_scale,
+                "sampler_name": "euler",
+                "scheduler": "simple",
+                "denoise": 1.0
+            }
+        }
+        node_id += 1
+
+        # --- Decode + output ---
+
+        decode_id = str(node_id)
+        workflow[decode_id] = {
+            "class_type": "VAEDecode",
+            "inputs": {
+                "samples": [sampler_id, 0],
+                "vae": [vae_id, 0]
+            }
+        }
+        node_id += 1
+
+        preview_id = str(node_id)
+        workflow[preview_id] = {
+            "class_type": "PreviewImage",
+            "inputs": {
+                "images": [decode_id, 0]
+            }
+        }
+
+        return workflow
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        num_inference_steps: Optional[int] = None,
+        cfg_scale: Optional[float] = None
+    ) -> GenerationResult:
+        """
+        Generate/edit image using Qwen-Image-Edit-2511 via ComfyUI.
+
+        Generates at native 1024x1024, then crop+resize to requested
+        aspect ratio for clean VAE output.
+        """
+        if not self.is_healthy():
+            return GenerationResult.error_result(
+                "ComfyUI server is not accessible. Make sure ComfyUI is running on "
+                f"{self.server_address}"
+            )
+
+        try:
+            start_time = time.time()
+
+            # Target dimensions for post-processing
+            target_w, target_h = self._get_dimensions(request.aspect_ratio)
+            # Generate at native resolution (VAE-safe)
+            native_w, native_h = self.NATIVE_RESOLUTION
+
+            # Upload input images (max 3)
+            uploaded_images = []
+            if request.has_input_images:
+                for i, img in enumerate(request.input_images):
+                    if img is not None:
+                        name = f"input_{i}_{uuid.uuid4().hex[:8]}.png"
+                        uploaded_name = self._upload_image(img, name)
+                        if uploaded_name:
+                            uploaded_images.append(uploaded_name)
+                        else:
+                            logger.warning(f"Failed to upload image {i}")
+
+            steps = num_inference_steps or self.num_inference_steps
+            cfg = cfg_scale or self.cfg_scale
+
+            # Temporarily set for workflow build
+            old_steps, old_cfg = self.num_inference_steps, self.cfg_scale
+            self.num_inference_steps, self.cfg_scale = steps, cfg
+
+            workflow = self._build_workflow(
+                prompt=request.prompt,
+                width=native_w,
+                height=native_h,
+                input_images=uploaded_images or None,
+                negative_prompt=request.negative_prompt or ""
+            )
+
+            self.num_inference_steps, self.cfg_scale = old_steps, old_cfg
+
+            logger.info(f"Generating with ComfyUI/Qwen: {request.prompt[:80]}...")
+            logger.info(
+                f"Native: {native_w}x{native_h}, target: {target_w}x{target_h}, "
+                f"steps: {steps}, cfg: {cfg}, images: {len(uploaded_images)}, "
+                f"lightning: {self.use_lightning}"
+            )
+
+            # Queue and wait
+            prompt_id = self._queue_prompt(workflow)
+            logger.info(f"Queued prompt: {prompt_id}")
+
+            if not self._wait_for_completion(prompt_id):
+                return GenerationResult.error_result("Generation failed or timed out")
+
+            # Retrieve output
+            history = self._get_history(prompt_id)
+            if prompt_id not in history:
+                return GenerationResult.error_result("No history found for prompt")
+
+            outputs = history[prompt_id].get('outputs', {})
+            for nid, node_output in outputs.items():
+                if 'images' in node_output:
+                    for img_info in node_output['images']:
+                        img_data = self._get_image(
+                            img_info['filename'],
+                            img_info.get('subfolder', ''),
+                            img_info.get('type', 'temp')
+                        )
+                        image = Image.open(io.BytesIO(img_data))
+                        generation_time = time.time() - start_time
+                        logger.info(f"Generated in {generation_time:.2f}s: {image.size}")
+
+                        # Crop+resize to target aspect ratio
+                        if (target_w, target_h) != (native_w, native_h):
+                            image = self._crop_and_resize(image, target_w, target_h)
+                            logger.info(f"Post-processed to: {image.size}")
+
+                        return GenerationResult.success_result(
+                            image=image,
+                            message=f"Generated with ComfyUI/Qwen in {generation_time:.2f}s",
+                            generation_time=generation_time
+                        )
+
+            return GenerationResult.error_result("No output images found")
+
+        except Exception as e:
+            logger.error(f"ComfyUI generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"ComfyUI error: {str(e)}")
+
+    def unload_model(self):
+        """
+        Request ComfyUI to free memory.
+        Note: ComfyUI manages models automatically, but we can request cleanup.
+        """
+        try:
+            url = f"http://{self.server_address}/free"
+            data = json.dumps({"unload_models": True}).encode('utf-8')
+            req = urllib.request.Request(url, data=data, method='POST')
+            urllib.request.urlopen(req)
+            logger.info("Requested ComfyUI to free memory")
+        except Exception as e:
+            logger.warning(f"Failed to request memory cleanup: {e}")
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> Tuple[int, int]:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))

src/core/__init__.py

@@ -0,0 +1,16 @@
+"""
+Character Sheet Pro - Core API
+==============================
+
+Exposes core functionality for plugins.
+"""
+
+from .plugin_base import Plugin, PluginMetadata, PluginAPI
+from .plugin_manager import PluginManager
+
+__all__ = [
+    'Plugin',
+    'PluginMetadata',
+    'PluginAPI',
+    'PluginManager'
+]

src/flux_klein_client.py

@@ -0,0 +1,270 @@
+"""
+FLUX.2 Klein Client
+===================
+
+Client for FLUX.2 klein 4B local image generation.
+Supports text-to-image and multi-reference editing.
+"""
+
+import logging
+import time
+from typing import Optional, List
+from PIL import Image
+
+import torch
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class FluxKleinClient:
+    """
+    Client for FLUX.2 klein models.
+
+    Supports:
+    - Text-to-image generation
+    - Single and multi-reference image editing
+    - Multiple model sizes (4B, 9B) and variants (distilled, base)
+    """
+
+    # Model variants - choose based on quality/speed tradeoff
+    MODELS = {
+        # 4B models (~13GB VRAM)
+        "4b": "black-forest-labs/FLUX.2-klein-4B",           # Distilled, 4 steps
+        "4b-base": "black-forest-labs/FLUX.2-klein-base-4B", # Base, configurable steps
+        # 9B models (~29GB VRAM, better quality)
+        "9b": "black-forest-labs/FLUX.2-klein-9B",           # Distilled, 4 steps
+        "9b-base": "black-forest-labs/FLUX.2-klein-base-9B", # Base, 50 steps - BEST QUALITY
+        "9b-fp8": "black-forest-labs/FLUX.2-klein-9b-fp8",   # FP8 quantized (~20GB)
+    }
+
+    # Legacy compatibility
+    MODEL_ID = MODELS["4b"]
+    MODEL_ID_BASE = MODELS["4b-base"]
+
+    # Aspect ratio to dimensions mapping
+    ASPECT_RATIOS = {
+        "1:1": (1024, 1024),
+        "16:9": (1344, 768),
+        "9:16": (768, 1344),
+        "21:9": (1536, 640),    # Cinematic ultra-wide
+        "3:2": (1248, 832),
+        "2:3": (832, 1248),
+        "3:4": (896, 1152),
+        "4:3": (1152, 896),
+        "4:5": (896, 1120),
+        "5:4": (1120, 896),
+    }
+
+    # Default settings for each model variant
+    MODEL_DEFAULTS = {
+        "4b": {"steps": 4, "guidance": 1.0},
+        "4b-base": {"steps": 28, "guidance": 3.5},
+        "9b": {"steps": 4, "guidance": 1.0},
+        "9b-base": {"steps": 50, "guidance": 4.0},  # Best quality
+        "9b-fp8": {"steps": 4, "guidance": 4.0},
+    }
+
+    def __init__(
+        self,
+        model_variant: str = "9b-base",  # Default to highest quality
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        enable_cpu_offload: bool = True,
+        # Legacy params
+        use_base_model: bool = False,
+    ):
+        """
+        Initialize FLUX.2 klein client.
+
+        Args:
+            model_variant: Model variant to use:
+                - "4b": Fast, 4 steps, ~13GB VRAM
+                - "4b-base": Configurable steps, ~13GB VRAM
+                - "9b": Better quality, 4 steps, ~29GB VRAM
+                - "9b-base": BEST quality, 50 steps, ~29GB VRAM
+                - "9b-fp8": FP8 quantized, ~20GB VRAM
+            device: Device to use (cuda or cpu)
+            dtype: Data type for model weights
+            enable_cpu_offload: Enable CPU offload to save VRAM
+        """
+        # Handle legacy use_base_model parameter
+        if use_base_model and model_variant == "9b-base":
+            model_variant = "4b-base"
+
+        self.model_variant = model_variant
+        self.device = device
+        self.dtype = dtype
+        self.enable_cpu_offload = enable_cpu_offload
+        self.pipe = None
+        self._loaded = False
+
+        # Get default settings for this variant
+        defaults = self.MODEL_DEFAULTS.get(model_variant, {"steps": 4, "guidance": 1.0})
+        self.default_steps = defaults["steps"]
+        self.default_guidance = defaults["guidance"]
+
+        logger.info(f"FluxKleinClient initialized (variant: {model_variant}, steps: {self.default_steps}, guidance: {self.default_guidance})")
+
+    def load_model(self) -> bool:
+        """Load the model into memory."""
+        if self._loaded:
+            return True
+
+        try:
+            # Get model ID for selected variant
+            model_id = self.MODELS.get(self.model_variant, self.MODELS["4b"])
+            logger.info(f"Loading FLUX.2 klein ({self.model_variant}) from {model_id}...")
+
+            start_time = time.time()
+
+            # FLUX.2 klein requires Flux2KleinPipeline (specific to klein models)
+            # Requires diffusers from git: pip install git+https://github.com/huggingface/diffusers.git
+            from diffusers import Flux2KleinPipeline
+
+            self.pipe = Flux2KleinPipeline.from_pretrained(
+                model_id,
+                torch_dtype=self.dtype,
+            )
+
+            # Use enable_model_cpu_offload() for VRAM management (documented approach)
+            if self.enable_cpu_offload:
+                self.pipe.enable_model_cpu_offload()
+                logger.info("CPU offload enabled")
+            else:
+                self.pipe.to(self.device)
+                logger.info(f"Model moved to {self.device}")
+
+            load_time = time.time() - start_time
+            logger.info(f"FLUX.2 klein ({self.model_variant}) loaded in {load_time:.1f}s")
+
+            # Validate by running a test generation
+            logger.info("Validating model with test generation...")
+            try:
+                test_result = self.pipe(
+                    prompt="A simple test image",
+                    height=256,
+                    width=256,
+                    guidance_scale=1.0,
+                    num_inference_steps=1,
+                    generator=torch.Generator(device="cpu").manual_seed(42),
+                )
+                if test_result.images[0] is not None:
+                    logger.info("Model validation successful")
+                else:
+                    logger.error("Model validation failed: no output image")
+                    return False
+            except Exception as e:
+                logger.error(f"Model validation failed: {e}", exc_info=True)
+                return False
+
+            self._loaded = True
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to load FLUX.2 klein: {e}", exc_info=True)
+            return False
+
+    def unload_model(self):
+        """Unload model from memory."""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+            self._loaded = False
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            logger.info("FLUX.2 klein unloaded")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        num_inference_steps: int = None,
+        guidance_scale: float = None
+    ) -> GenerationResult:
+        """
+        Generate image using FLUX.2 klein.
+
+        Args:
+            request: GenerationRequest object
+            num_inference_steps: Number of denoising steps (4 for klein distilled)
+            guidance_scale: Classifier-free guidance scale
+
+        Returns:
+            GenerationResult object
+        """
+        if not self._loaded:
+            if not self.load_model():
+                return GenerationResult.error_result("Failed to load FLUX.2 klein model")
+
+        # Use model defaults if not specified
+        if num_inference_steps is None:
+            num_inference_steps = self.default_steps
+        if guidance_scale is None:
+            guidance_scale = self.default_guidance
+
+        try:
+            start_time = time.time()
+
+            # Get dimensions from aspect ratio
+            width, height = self._get_dimensions(request.aspect_ratio)
+
+            logger.info(f"Generating with {self.model_variant}: steps={num_inference_steps}, guidance={guidance_scale}")
+
+            # Build generation kwargs
+            gen_kwargs = {
+                "prompt": request.prompt,
+                "height": height,
+                "width": width,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps,
+                "generator": torch.Generator(device="cpu").manual_seed(42),
+            }
+
+            # Add input images if present (for editing)
+            if request.has_input_images:
+                # FLUX.2 klein supports multi-reference editing
+                # Pass images as 'image' parameter
+                valid_images = [img for img in request.input_images if img is not None]
+                if len(valid_images) == 1:
+                    gen_kwargs["image"] = valid_images[0]
+                elif len(valid_images) > 1:
+                    gen_kwargs["image"] = valid_images
+
+            logger.info(f"Generating with FLUX.2 klein: {request.prompt[:80]}...")
+
+            # Generate
+            with torch.inference_mode():
+                output = self.pipe(**gen_kwargs)
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Generated in {generation_time:.2f}s: {image.size}")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Generated with FLUX.2 klein in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"FLUX.2 klein generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"FLUX.2 klein error: {str(e)}")
+
+    def _get_dimensions(self, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return self.ASPECT_RATIOS.get(ratio, (1024, 1024))
+
+    def is_healthy(self) -> bool:
+        """Check if model is loaded and ready."""
+        return self._loaded and self.pipe is not None
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))

src/gemini_client.py

@@ -0,0 +1,224 @@
+"""
+Gemini API Client
+=================
+
+Client for Google Gemini Image APIs (Flash and Pro models).
+Handles API communication and response parsing.
+"""
+
+import base64
+import logging
+from io import BytesIO
+from typing import Optional
+from PIL import Image
+
+from google import genai
+from google.genai import types
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiClient:
+    """
+    Client for Gemini Image APIs.
+
+    Supports:
+    - Gemini 2.5 Flash Image (up to ~3 reference images)
+    - Gemini 3 Pro Image Preview (up to 14 reference images, 1K/2K/4K)
+    """
+
+    # Model names (updated January 2026)
+    # See: https://ai.google.dev/gemini-api/docs/image-generation
+    MODEL_FLASH = "gemini-2.5-flash-image"      # Fast, efficient image generation
+    MODEL_PRO = "gemini-3-pro-image-preview"    # Pro quality, advanced text rendering
+
+    # Valid resolutions for Pro model
+    VALID_RESOLUTIONS = ["1K", "2K", "4K"]
+
+    # Aspect ratio to dimensions mapping
+    ASPECT_RATIOS = {
+        "1:1": (1024, 1024),
+        "16:9": (1344, 768),
+        "9:16": (768, 1344),
+        "21:9": (1536, 640),    # Cinematic ultra-wide
+        "3:2": (1248, 832),
+        "2:3": (832, 1248),
+        "3:4": (864, 1184),
+        "4:3": (1344, 1008),
+        "4:5": (1024, 1280),
+        "5:4": (1280, 1024),
+    }
+
+    def __init__(self, api_key: str, use_pro_model: bool = False):
+        """
+        Initialize Gemini client.
+
+        Args:
+            api_key: Google Gemini API key
+            use_pro_model: If True, use Pro model with enhanced capabilities
+        """
+        if not api_key:
+            raise ValueError("API key is required for Gemini client")
+
+        self.api_key = api_key
+        self.use_pro_model = use_pro_model
+        self.client = genai.Client(api_key=api_key)
+
+        model_name = self.MODEL_PRO if use_pro_model else self.MODEL_FLASH
+        logger.info(f"GeminiClient initialized with model: {model_name}")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        resolution: str = "1K"
+    ) -> GenerationResult:
+        """
+        Generate image using Gemini API.
+
+        Args:
+            request: GenerationRequest object
+            resolution: Resolution for Pro model ("1K", "2K", "4K")
+
+        Returns:
+            GenerationResult object
+        """
+        try:
+            model_name = self.MODEL_PRO if self.use_pro_model else self.MODEL_FLASH
+            logger.info(f"Generating with {model_name}: {request.prompt[:100]}...")
+
+            # Build contents list
+            contents = self._build_contents(request)
+
+            # Build config
+            config = self._build_config(
+                request,
+                resolution if self.use_pro_model else None
+            )
+
+            # Call API
+            response = self.client.models.generate_content(
+                model=model_name,
+                contents=contents,
+                config=config
+            )
+
+            # Parse response
+            return self._parse_response(response)
+
+        except Exception as e:
+            logger.error(f"Gemini generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Gemini API error: {str(e)}")
+
+    def _build_contents(self, request: GenerationRequest) -> list:
+        """Build contents list for API request."""
+        contents = []
+
+        # Add input images if present
+        if request.has_input_images:
+            valid_images = [img for img in request.input_images if img is not None]
+            contents.extend(valid_images)
+
+        # Add prompt
+        contents.append(request.prompt)
+
+        return contents
+
+    def _build_config(
+        self,
+        request: GenerationRequest,
+        resolution: Optional[str] = None
+    ) -> types.GenerateContentConfig:
+        """Build generation config for API request."""
+        # Parse aspect ratio
+        aspect_ratio = request.aspect_ratio
+        if " " in aspect_ratio:
+            aspect_ratio = aspect_ratio.split()[0]
+
+        # Build image config
+        image_config_kwargs = {"aspect_ratio": aspect_ratio}
+
+        # Add resolution for Pro model
+        if resolution and self.use_pro_model:
+            if resolution not in self.VALID_RESOLUTIONS:
+                logger.warning(f"Invalid resolution '{resolution}', defaulting to '1K'")
+                resolution = "1K"
+            image_config_kwargs["output_image_resolution"] = resolution
+            logger.info(f"Pro model resolution: {resolution}")
+
+        config = types.GenerateContentConfig(
+            temperature=request.temperature,
+            response_modalities=["image", "text"],
+            image_config=types.ImageConfig(**image_config_kwargs)
+        )
+
+        return config
+
+    def _parse_response(self, response) -> GenerationResult:
+        """Parse API response and extract image."""
+        if response is None:
+            return GenerationResult.error_result("No response from API")
+
+        if not hasattr(response, 'candidates') or not response.candidates:
+            return GenerationResult.error_result("No candidates in response")
+
+        candidate = response.candidates[0]
+
+        # Check finish reason
+        if hasattr(candidate, 'finish_reason'):
+            finish_reason = str(candidate.finish_reason)
+            logger.info(f"Finish reason: {finish_reason}")
+
+            if 'SAFETY' in finish_reason or 'PROHIBITED' in finish_reason:
+                return GenerationResult.error_result(
+                    f"Content blocked by safety filters: {finish_reason}"
+                )
+
+        # Check for content
+        if not hasattr(candidate, 'content') or candidate.content is None:
+            finish_reason = getattr(candidate, 'finish_reason', 'UNKNOWN')
+            return GenerationResult.error_result(
+                f"No content in response (finish_reason: {finish_reason})"
+            )
+
+        # Extract image from parts
+        if hasattr(candidate.content, 'parts') and candidate.content.parts:
+            for part in candidate.content.parts:
+                if hasattr(part, 'inline_data') and part.inline_data:
+                    try:
+                        image_data = part.inline_data.data
+
+                        # Handle both bytes and base64 string
+                        if isinstance(image_data, str):
+                            image_data = base64.b64decode(image_data)
+
+                        # Convert to PIL Image
+                        image_buffer = BytesIO(image_data)
+                        image = Image.open(image_buffer)
+                        image.load()
+
+                        logger.info(f"Image generated: {image.size}, {image.mode}")
+                        return GenerationResult.success_result(
+                            image=image,
+                            message="Generated successfully"
+                        )
+
+                    except Exception as e:
+                        logger.error(f"Failed to decode image: {e}")
+                        return GenerationResult.error_result(
+                            f"Image decoding error: {str(e)}"
+                        )
+
+        return GenerationResult.error_result("No image data in response")
+
+    def is_healthy(self) -> bool:
+        """Check if API is accessible."""
+        return self.api_key is not None and len(self.api_key) > 0
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))

src/longcat_edit_client.py

@@ -0,0 +1,301 @@
+"""
+LongCat-Image-Edit Client
+=========================
+
+Client for Meituan's LongCat-Image-Edit model.
+Supports instruction-following image editing with bilingual (Chinese-English) support.
+
+This is a SOTA open-source image editing model with excellent:
+- Global editing, local editing, text modification
+- Reference-guided editing
+- Consistency preservation (layout, texture, color tone, identity)
+- Multi-turn editing capabilities
+"""
+
+import logging
+import time
+from typing import Optional, List
+from PIL import Image
+
+import torch
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class LongCatEditClient:
+    """
+    Client for LongCat-Image-Edit model from Meituan.
+
+    Features:
+    - Instruction-following image editing
+    - Bilingual support (Chinese-English)
+    - Excellent consistency preservation
+    - Multi-turn editing
+
+    Requires ~18GB VRAM with CPU offload.
+    """
+
+    MODEL_ID = "meituan-longcat/LongCat-Image-Edit"
+
+    # Aspect ratio to dimensions mapping
+    ASPECT_RATIOS = {
+        "1:1": (1024, 1024),
+        "16:9": (1344, 768),
+        "9:16": (768, 1344),
+        "21:9": (1536, 640),    # Cinematic ultra-wide
+        "3:2": (1248, 832),
+        "2:3": (832, 1248),
+        "3:4": (896, 1152),
+        "4:3": (1152, 896),
+        "4:5": (896, 1120),
+        "5:4": (1120, 896),
+    }
+
+    # Default generation settings
+    DEFAULT_STEPS = 50
+    DEFAULT_GUIDANCE = 4.5
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        enable_cpu_offload: bool = True,
+    ):
+        """
+        Initialize LongCat-Image-Edit client.
+
+        Args:
+            device: Device to use (cuda or cpu)
+            dtype: Data type for model weights (bfloat16 recommended)
+            enable_cpu_offload: Enable CPU offload to save VRAM (~18GB required)
+        """
+        self.device = device
+        self.dtype = dtype
+        self.enable_cpu_offload = enable_cpu_offload
+        self.pipe = None
+        self._loaded = False
+
+        logger.info(f"LongCatEditClient initialized (cpu_offload: {enable_cpu_offload})")
+
+    def load_model(self) -> bool:
+        """Load the model into memory."""
+        if self._loaded:
+            return True
+
+        try:
+            logger.info(f"Loading LongCat-Image-Edit from {self.MODEL_ID}...")
+
+            start_time = time.time()
+
+            # Import LongCat pipeline
+            # Requires latest diffusers: pip install git+https://github.com/huggingface/diffusers
+            from diffusers import LongCatImageEditPipeline
+
+            self.pipe = LongCatImageEditPipeline.from_pretrained(
+                self.MODEL_ID,
+                torch_dtype=self.dtype,
+            )
+
+            # Apply memory optimization
+            if self.enable_cpu_offload:
+                self.pipe.enable_model_cpu_offload()
+                logger.info("CPU offload enabled (~18GB VRAM)")
+            else:
+                self.pipe.to(self.device, self.dtype)
+                logger.info(f"Model moved to {self.device} (high VRAM mode)")
+
+            load_time = time.time() - start_time
+            logger.info(f"LongCat-Image-Edit loaded in {load_time:.1f}s")
+
+            self._loaded = True
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to load LongCat-Image-Edit: {e}", exc_info=True)
+            return False
+
+    def unload_model(self):
+        """Unload model from memory."""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+
+        self._loaded = False
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        logger.info("LongCat-Image-Edit unloaded")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        num_inference_steps: int = None,
+        guidance_scale: float = None
+    ) -> GenerationResult:
+        """
+        Edit image using LongCat-Image-Edit.
+
+        Args:
+            request: GenerationRequest object with:
+                - prompt: The editing instruction (e.g., "Change the background to a forest")
+                - input_images: List with the source image to edit
+                - aspect_ratio: Output aspect ratio
+            num_inference_steps: Number of denoising steps (default: 50)
+            guidance_scale: Classifier-free guidance scale (default: 4.5)
+
+        Returns:
+            GenerationResult object
+        """
+        if not self._loaded:
+            if not self.load_model():
+                return GenerationResult.error_result("Failed to load LongCat-Image-Edit model")
+
+        # Use defaults if not specified
+        if num_inference_steps is None:
+            num_inference_steps = self.DEFAULT_STEPS
+        if guidance_scale is None:
+            guidance_scale = self.DEFAULT_GUIDANCE
+
+        try:
+            start_time = time.time()
+
+            # Get input image
+            if not request.has_input_images:
+                return GenerationResult.error_result("LongCat-Image-Edit requires an input image to edit")
+
+            input_image = None
+            for img in request.input_images:
+                if img is not None:
+                    input_image = img
+                    break
+
+            if input_image is None:
+                return GenerationResult.error_result("No valid input image provided")
+
+            # Get dimensions from aspect ratio
+            width, height = self._get_dimensions(request.aspect_ratio)
+
+            # Resize input image to target dimensions
+            input_image = input_image.convert('RGB')
+            input_image = input_image.resize((width, height), Image.Resampling.LANCZOS)
+
+            logger.info(f"Editing with LongCat: steps={num_inference_steps}, guidance={guidance_scale}")
+            logger.info(f"Edit instruction: {request.prompt[:100]}...")
+
+            # Build generation kwargs
+            gen_kwargs = {
+                "image": input_image,
+                "prompt": request.prompt,
+                "negative_prompt": request.negative_prompt or "",
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps,
+                "num_images_per_prompt": 1,
+                "generator": torch.Generator("cpu").manual_seed(42),
+            }
+
+            # Generate
+            with torch.inference_mode():
+                output = self.pipe(**gen_kwargs)
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Edited in {generation_time:.2f}s: {image.size}")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Edited with LongCat-Image-Edit in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"LongCat-Image-Edit generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"LongCat-Image-Edit error: {str(e)}")
+
+    def edit_with_instruction(
+        self,
+        source_image: Image.Image,
+        instruction: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = None,
+        guidance_scale: float = None,
+        seed: int = 42
+    ) -> GenerationResult:
+        """
+        Simplified method for instruction-based image editing.
+
+        Args:
+            source_image: The image to edit
+            instruction: Natural language editing instruction
+                Examples:
+                - "Change the background to a sunset beach"
+                - "Make the person wear a red dress"
+                - "Add snow to the scene"
+                - "Change the cat to a dog"
+            negative_prompt: What to avoid in the output
+            num_inference_steps: Denoising steps (default: 50)
+            guidance_scale: CFG scale (default: 4.5)
+            seed: Random seed for reproducibility
+
+        Returns:
+            GenerationResult with the edited image
+        """
+        if not self._loaded:
+            if not self.load_model():
+                return GenerationResult.error_result("Failed to load LongCat-Image-Edit model")
+
+        if num_inference_steps is None:
+            num_inference_steps = self.DEFAULT_STEPS
+        if guidance_scale is None:
+            guidance_scale = self.DEFAULT_GUIDANCE
+
+        try:
+            start_time = time.time()
+
+            # Ensure RGB
+            source_image = source_image.convert('RGB')
+
+            logger.info(f"Editing image with instruction: {instruction[:100]}...")
+
+            with torch.inference_mode():
+                output = self.pipe(
+                    image=source_image,
+                    prompt=instruction,
+                    negative_prompt=negative_prompt,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    num_images_per_prompt=1,
+                    generator=torch.Generator("cpu").manual_seed(seed),
+                )
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Edit completed in {generation_time:.2f}s")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Edited with instruction in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"Instruction-based edit failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Edit error: {str(e)}")
+
+    def _get_dimensions(self, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return self.ASPECT_RATIOS.get(ratio, (1024, 1024))
+
+    def is_healthy(self) -> bool:
+        """Check if model is loaded and ready."""
+        return self._loaded and self.pipe is not None
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))

src/model_manager.py

@@ -0,0 +1,332 @@
+"""
+Model Manager
+=============
+
+Manages model loading states and provides a robust interface for
+ensuring models are loaded and validated before generation.
+
+States:
+- UNLOADED: No model loaded
+- LOADING: Model is being loaded
+- READY: Model loaded and validated
+- ERROR: Model failed to load
+"""
+
+import logging
+import threading
+import time
+from typing import Optional, Callable, Tuple
+from enum import Enum
+from PIL import Image
+
+from .backend_router import BackendRouter, BackendType
+from .character_service import CharacterSheetService
+
+
+logger = logging.getLogger(__name__)
+
+
+class ModelState(Enum):
+    """Model loading states."""
+    UNLOADED = "unloaded"
+    LOADING = "loading"
+    READY = "ready"
+    ERROR = "error"
+
+
+class ModelManager:
+    """
+    Manages model loading lifecycle with state tracking.
+
+    Ensures models are fully loaded and validated before allowing generation.
+    Provides progress callbacks for UI updates during loading.
+    """
+
+    def __init__(self):
+        self._state = ModelState.UNLOADED
+        self._current_backend: Optional[BackendType] = None
+        self._service: Optional[CharacterSheetService] = None
+        self._error_message: Optional[str] = None
+        self._loading_progress: float = 0.0
+        self._loading_message: str = ""
+        self._lock = threading.Lock()
+        self._cancel_requested = False
+
+    @property
+    def state(self) -> ModelState:
+        """Current model state."""
+        return self._state
+
+    @property
+    def is_ready(self) -> bool:
+        """Check if model is ready for generation."""
+        return self._state == ModelState.READY
+
+    @property
+    def is_loading(self) -> bool:
+        """Check if model is currently loading."""
+        return self._state == ModelState.LOADING
+
+    @property
+    def error_message(self) -> Optional[str]:
+        """Get error message if in error state."""
+        return self._error_message
+
+    @property
+    def loading_progress(self) -> float:
+        """Get loading progress (0.0 to 1.0)."""
+        return self._loading_progress
+
+    @property
+    def loading_message(self) -> str:
+        """Get current loading status message."""
+        return self._loading_message
+
+    @property
+    def current_backend(self) -> Optional[BackendType]:
+        """Get currently loaded backend."""
+        return self._current_backend
+
+    @property
+    def service(self) -> Optional[CharacterSheetService]:
+        """Get the character sheet service (only valid when ready)."""
+        if self._state != ModelState.READY:
+            return None
+        return self._service
+
+    def get_status_display(self) -> Tuple[str, str]:
+        """
+        Get status message and color for UI display.
+
+        Returns:
+            Tuple of (message, color) where color is a CSS color string
+        """
+        if self._state == ModelState.UNLOADED:
+            return "No model loaded", "#888888"
+        elif self._state == ModelState.LOADING:
+            pct = int(self._loading_progress * 100)
+            return f"Loading... {pct}% - {self._loading_message}", "#FFA500"
+        elif self._state == ModelState.READY:
+            backend_name = BackendRouter.BACKEND_NAMES.get(
+                self._current_backend,
+                str(self._current_backend)
+            )
+            return f"Ready: {backend_name}", "#00AA00"
+        elif self._state == ModelState.ERROR:
+            return f"Error: {self._error_message}", "#FF0000"
+        return "Unknown state", "#888888"
+
+    def request_cancel(self):
+        """Request cancellation of current loading operation."""
+        self._cancel_requested = True
+        logger.info("Model loading cancellation requested")
+
+    def load_model(
+        self,
+        backend: BackendType,
+        api_key: Optional[str] = None,
+        steps: int = 4,
+        guidance: float = 1.0,
+        progress_callback: Optional[Callable[[float, str], None]] = None
+    ) -> bool:
+        """
+        Load a model with progress tracking.
+
+        Args:
+            backend: Backend type to load
+            api_key: API key for cloud backends
+            steps: Default steps for generation
+            guidance: Default guidance scale
+            progress_callback: Callback for progress updates (progress, message)
+
+        Returns:
+            True if model loaded successfully
+        """
+        with self._lock:
+            if self._state == ModelState.LOADING:
+                logger.warning("Model is already loading, ignoring request")
+                return False
+
+            self._state = ModelState.LOADING
+            self._loading_progress = 0.0
+            self._loading_message = "Initializing..."
+            self._error_message = None
+            self._cancel_requested = False
+
+        def update_progress(progress: float, message: str):
+            self._loading_progress = progress
+            self._loading_message = message
+            if progress_callback:
+                progress_callback(progress, message)
+
+        try:
+            # Step 1: Unload previous model if different backend
+            update_progress(0.05, "Checking current model...")
+
+            if self._service and self._current_backend != backend:
+                update_progress(0.1, "Unloading previous model...")
+                try:
+                    if hasattr(self._service, 'router'):
+                        self._service.router.unload_local_models()
+                except Exception as e:
+                    logger.warning(f"Error unloading previous model: {e}")
+                self._service = None
+
+            if self._cancel_requested:
+                self._state = ModelState.UNLOADED
+                return False
+
+            # Step 2: Create service and load model
+            backend_name = BackendRouter.BACKEND_NAMES.get(backend, str(backend))
+            update_progress(0.15, f"Loading {backend_name}...")
+
+            logger.info(f"Creating CharacterSheetService for {backend.value}")
+
+            # For local models, this will load the model
+            # For cloud backends, this just validates the API key
+            self._service = CharacterSheetService(
+                api_key=api_key,
+                backend=backend
+            )
+
+            if self._cancel_requested:
+                self._state = ModelState.UNLOADED
+                self._service = None
+                return False
+
+            update_progress(0.7, "Model loaded, configuring...")
+
+            # Step 3: Configure default parameters
+            if hasattr(self._service.client, 'default_steps'):
+                self._service.client.default_steps = steps
+            if hasattr(self._service.client, 'default_guidance'):
+                self._service.client.default_guidance = guidance
+
+            update_progress(0.8, "Validating model...")
+
+            # Step 4: Validate model is actually working
+            is_valid, error = self._validate_model()
+
+            if not is_valid:
+                raise RuntimeError(f"Model validation failed: {error}")
+
+            update_progress(1.0, "Ready!")
+
+            # Success!
+            with self._lock:
+                self._current_backend = backend
+                self._state = ModelState.READY
+                self._loading_progress = 1.0
+                self._loading_message = "Ready"
+
+            logger.info(f"Model {backend.value} loaded and validated successfully")
+            return True
+
+        except Exception as e:
+            error_msg = str(e)
+            logger.error(f"Failed to load model {backend.value}: {error_msg}", exc_info=True)
+
+            with self._lock:
+                self._state = ModelState.ERROR
+                self._error_message = self._simplify_error(error_msg)
+                self._service = None
+
+            if progress_callback:
+                progress_callback(0.0, f"Error: {self._error_message}")
+
+            return False
+
+    def _validate_model(self) -> Tuple[bool, Optional[str]]:
+        """
+        Validate that the model is actually working.
+
+        For local models, checks that the pipeline is loaded.
+        For cloud backends, does a minimal health check.
+
+        Returns:
+            Tuple of (is_valid, error_message)
+        """
+        if self._service is None:
+            return False, "Service not initialized"
+
+        try:
+            client = self._service.client
+
+            # Check if client has health check method
+            if hasattr(client, 'is_healthy'):
+                if not client.is_healthy():
+                    return False, "Client health check failed"
+
+            # For local models, check pipeline is loaded
+            if hasattr(client, '_loaded'):
+                if not client._loaded:
+                    return False, "Model pipeline not loaded"
+
+            # For FLUX models, verify the pipe exists
+            if hasattr(client, 'pipe'):
+                if client.pipe is None:
+                    return False, "Model pipeline is None"
+
+            return True, None
+
+        except Exception as e:
+            return False, str(e)
+
+    def _simplify_error(self, error: str) -> str:
+        """Simplify technical error messages for user display."""
+        error_lower = error.lower()
+
+        if "cuda out of memory" in error_lower or "out of memory" in error_lower:
+            return "Not enough GPU memory. Try a smaller model or close other applications."
+
+        if "api key" in error_lower:
+            return "Invalid or missing API key."
+
+        if "connection" in error_lower or "network" in error_lower:
+            return "Network connection error. Check your internet connection."
+
+        if "not found" in error_lower and "model" in error_lower:
+            return "Model files not found. The model may need to be downloaded."
+
+        if "import" in error_lower:
+            return "Missing dependencies. Some required packages are not installed."
+
+        if "meta tensor" in error_lower:
+            return "Model loading failed (meta tensor error). Try restarting the application."
+
+        # Truncate long errors
+        if len(error) > 100:
+            return error[:97] + "..."
+
+        return error
+
+    def unload(self):
+        """Unload the current model."""
+        with self._lock:
+            if self._service:
+                try:
+                    if hasattr(self._service, 'router'):
+                        self._service.router.unload_local_models()
+                except Exception as e:
+                    logger.warning(f"Error during unload: {e}")
+                self._service = None
+
+            self._state = ModelState.UNLOADED
+            self._current_backend = None
+            self._error_message = None
+            self._loading_progress = 0.0
+            self._loading_message = ""
+
+        logger.info("Model unloaded")
+
+
+# Global singleton for model management
+_model_manager: Optional[ModelManager] = None
+
+
+def get_model_manager() -> ModelManager:
+    """Get the global ModelManager instance."""
+    global _model_manager
+    if _model_manager is None:
+        _model_manager = ModelManager()
+    return _model_manager

src/models.py

@@ -0,0 +1,94 @@
+"""
+Data Models for Character Sheet Pro
+====================================
+
+Dataclasses for generation requests and results.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, List, Dict, Any
+from PIL import Image
+from datetime import datetime
+
+
+@dataclass
+class GenerationRequest:
+    """Request for image generation."""
+
+    prompt: str
+    input_images: List[Image.Image] = field(default_factory=list)
+    aspect_ratio: str = "1:1"
+    temperature: float = 0.4
+    negative_prompt: Optional[str] = None
+
+    @property
+    def has_input_images(self) -> bool:
+        """Check if request has input images."""
+        return len(self.input_images) > 0
+
+
+@dataclass
+class GenerationResult:
+    """Result from image generation."""
+
+    success: bool
+    image: Optional[Image.Image] = None
+    message: str = ""
+    generation_time: Optional[float] = None
+
+    @classmethod
+    def success_result(
+        cls,
+        image: Image.Image,
+        message: str = "Generated successfully",
+        generation_time: Optional[float] = None
+    ) -> "GenerationResult":
+        """Create successful result."""
+        return cls(
+            success=True,
+            image=image,
+            message=message,
+            generation_time=generation_time
+        )
+
+    @classmethod
+    def error_result(cls, message: str) -> "GenerationResult":
+        """Create error result."""
+        return cls(success=False, message=message)
+
+
+@dataclass
+class CharacterSheetConfig:
+    """Configuration for character sheet generation."""
+
+    # Aspect ratios
+    face_aspect_ratio: str = "3:4"  # 864x1184
+    body_aspect_ratio: str = "9:16"  # 768x1344
+
+    # Generation temperatures
+    face_temperature: float = 0.35
+    body_temperature: float = 0.35
+    normalize_temperature: float = 0.5
+
+    # Layout
+    spacing: int = 20
+    background_color: str = "#2C2C2C"
+
+    # Retry settings
+    max_retries: int = 3
+    retry_delay: float = 30.0
+    rate_limit_delay_min: float = 2.0
+    rate_limit_delay_max: float = 3.0
+
+
+@dataclass
+class CharacterSheetMetadata:
+    """Metadata for generated character sheet."""
+
+    character_name: str
+    input_type: str  # "Face Only", "Full Body", "Face + Body"
+    costume_description: str
+    backend: str
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+    views: int = 7
+    stages: Dict[str, Any] = field(default_factory=dict)

src/qwen_image_edit_client.py

@@ -0,0 +1,495 @@
+"""
+Qwen-Image-Edit Client
+======================
+
+Client for Qwen-Image-Edit-2511 local image editing.
+Supports multi-image editing with improved consistency.
+
+GPU loading strategies (benchmarked on A6000 + A5000):
+  Pinned 2-GPU:        169.9s (4.25s/step) - 1.36x vs baseline
+  Balanced single-GPU: 184.4s (4.61s/step) - 1.25x vs baseline
+  CPU offload:         231.5s (5.79s/step) - baseline
+"""
+
+import logging
+import time
+import types
+from typing import Optional, List
+from PIL import Image
+
+import torch
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class QwenImageEditClient:
+    """
+    Client for Qwen-Image-Edit-2511 model.
+
+    Supports:
+    - Multi-image editing (up to multiple reference images)
+    - Precise text editing
+    - Improved character consistency
+    - LoRA integration
+    """
+
+    # Model variants
+    MODELS = {
+        "full": "Qwen/Qwen-Image-Edit",           # Official Qwen model
+    }
+
+    # Legacy compatibility
+    MODEL_ID = MODELS["full"]
+
+    # Aspect ratio to dimensions mapping (target output sizes)
+    ASPECT_RATIOS = {
+        "1:1": (1328, 1328),
+        "16:9": (1664, 928),
+        "9:16": (928, 1664),
+        "21:9": (1680, 720),    # Cinematic ultra-wide
+        "3:2": (1584, 1056),
+        "2:3": (1056, 1584),
+        "3:4": (1104, 1472),
+        "4:3": (1472, 1104),
+        "4:5": (1056, 1320),
+        "5:4": (1320, 1056),
+    }
+
+    # Proven native generation resolution.  Tested resolutions:
+    #   1104x1472 (3:4) → CLEAN output (face views in v1 test)
+    #   928x1664  (9:16) → VAE tiling noise / garbage
+    #   1328x1328 (1:1)  → VAE tiling noise / garbage
+    #   896x1184  (auto) → garbage
+    # Always generate at 1104x1472, then crop+resize to target.
+    NATIVE_RESOLUTION = (1104, 1472)
+
+    # VRAM thresholds for loading strategies
+    # Qwen-Image-Edit components: transformer ~40.9GB, text_encoder ~16.6GB, VAE ~0.25GB
+    BALANCED_VRAM_THRESHOLD_GB = 45    # Single GPU balanced (needs ~42GB + headroom)
+    MAIN_GPU_MIN_VRAM_GB = 42          # Transformer + VAE minimum
+    ENCODER_GPU_MIN_VRAM_GB = 17       # Text encoder minimum
+
+    def __init__(
+        self,
+        model_variant: str = "full",  # Use full model (~50GB)
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        enable_cpu_offload: bool = True,
+        encoder_device: Optional[str] = None,
+    ):
+        """
+        Initialize Qwen-Image-Edit client.
+
+        Args:
+            model_variant: Model variant ("full" for ~50GB)
+            device: Device to use for transformer+VAE (cuda or cuda:N)
+            dtype: Data type for model weights
+            enable_cpu_offload: Enable CPU offload to save VRAM
+            encoder_device: Explicit device for text_encoder (e.g. "cuda:3").
+                           If None, auto-detected from available GPUs.
+        """
+        self.model_variant = model_variant
+        self.device = device
+        self.dtype = dtype
+        self.enable_cpu_offload = enable_cpu_offload
+        self.encoder_device = encoder_device
+        self.pipe = None
+        self._loaded = False
+        self._loading_strategy = None
+
+        logger.info(f"QwenImageEditClient initialized (variant: {model_variant})")
+
+    @staticmethod
+    def _get_gpu_vram_gb(device_idx: int) -> float:
+        """Get total VRAM in GB for a specific GPU."""
+        if not torch.cuda.is_available():
+            return 0.0
+        if device_idx >= torch.cuda.device_count():
+            return 0.0
+        return torch.cuda.get_device_properties(device_idx).total_memory / 1e9
+
+    def _get_vram_gb(self) -> float:
+        """Get available VRAM in GB for the main target device."""
+        device_idx = self._parse_device_idx(self.device)
+        return self._get_gpu_vram_gb(device_idx)
+
+    @staticmethod
+    def _parse_device_idx(device: str) -> int:
+        """Parse CUDA device index from device string."""
+        if device.startswith("cuda:"):
+            try:
+                return int(device.split(":")[1])
+            except (ValueError, IndexError):
+                pass
+        return 0
+
+    def _find_encoder_gpu(self, main_idx: int) -> Optional[int]:
+        """Find a secondary GPU suitable for text_encoder (>= 17GB VRAM).
+
+        Prefers GPUs with more VRAM. Skips the main GPU.
+        """
+        if not torch.cuda.is_available():
+            return None
+
+        candidates = []
+        for i in range(torch.cuda.device_count()):
+            if i == main_idx:
+                continue
+            vram = self._get_gpu_vram_gb(i)
+            if vram >= self.ENCODER_GPU_MIN_VRAM_GB:
+                name = torch.cuda.get_device_name(i)
+                candidates.append((i, vram, name))
+
+        if not candidates:
+            return None
+
+        # Pick the GPU with the most VRAM
+        candidates.sort(key=lambda x: x[1], reverse=True)
+        best = candidates[0]
+        logger.info(f"Found encoder GPU: cuda:{best[0]} ({best[2]}, {best[1]:.1f} GB)")
+        return best[0]
+
+    @staticmethod
+    def _patched_get_qwen_prompt_embeds(self, prompt, image=None, device=None, dtype=None):
+        """Patched prompt encoding that routes inputs to text_encoder's device.
+
+        The original _get_qwen_prompt_embeds sends model_inputs to
+        execution_device (main GPU), then calls text_encoder on a different
+        GPU, causing a device mismatch. This patch:
+        1. Sends model_inputs to text_encoder's device for encoding
+        2. Moves outputs back to execution_device for the transformer
+        """
+        te_device = next(self.text_encoder.parameters()).device
+        execution_device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        template = self.prompt_template_encode
+        drop_idx = self.prompt_template_encode_start_idx
+        txt = [template.format(e) for e in prompt]
+
+        # Route to text_encoder's device, NOT execution_device
+        model_inputs = self.processor(
+            text=txt, images=image, padding=True, return_tensors="pt"
+        ).to(te_device)
+
+        outputs = self.text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=model_inputs.pixel_values,
+            image_grid_thw=model_inputs.image_grid_thw,
+            output_hidden_states=True,
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(
+            hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [
+            torch.ones(e.size(0), dtype=torch.long, device=e.device)
+            for e in split_hidden_states
+        ]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack([
+            torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
+            for u in split_hidden_states
+        ])
+        encoder_attention_mask = torch.stack([
+            torch.cat([u, u.new_zeros(max_seq_len - u.size(0))])
+            for u in attn_mask_list
+        ])
+
+        # Move outputs to execution_device for transformer
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=execution_device)
+        encoder_attention_mask = encoder_attention_mask.to(device=execution_device)
+
+        return prompt_embeds, encoder_attention_mask
+
+    def _load_pinned_multi_gpu(self, model_id: str, main_idx: int, encoder_idx: int) -> bool:
+        """Load with pinned multi-GPU: transformer+VAE on main, text_encoder on secondary.
+
+        Benchmarked at 169.9s (4.25s/step) - 1.36x faster than cpu_offload baseline.
+        """
+        from diffusers import QwenImageEditPipeline
+        from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+        from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformer2DModel
+        from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage
+        from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+
+        main_dev = f"cuda:{main_idx}"
+        enc_dev = f"cuda:{encoder_idx}"
+
+        logger.info(f"Loading pinned 2-GPU: transformer+VAE → {main_dev}, text_encoder → {enc_dev}")
+
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_id, subfolder="scheduler")
+        tokenizer = Qwen2Tokenizer.from_pretrained(
+            model_id, subfolder="tokenizer")
+        processor = Qwen2VLProcessor.from_pretrained(
+            model_id, subfolder="processor")
+
+        text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_id, subfolder="text_encoder", torch_dtype=self.dtype,
+        ).to(enc_dev)
+        logger.info(f"  text_encoder loaded on {enc_dev}")
+
+        transformer = QwenImageTransformer2DModel.from_pretrained(
+            model_id, subfolder="transformer", torch_dtype=self.dtype,
+        ).to(main_dev)
+        logger.info(f"  transformer loaded on {main_dev}")
+
+        vae = AutoencoderKLQwenImage.from_pretrained(
+            model_id, subfolder="vae", torch_dtype=self.dtype,
+        ).to(main_dev)
+        vae.enable_tiling()
+        logger.info(f"  VAE loaded on {main_dev}")
+
+        self.pipe = QwenImageEditPipeline(
+            scheduler=scheduler, vae=vae, text_encoder=text_encoder,
+            tokenizer=tokenizer, processor=processor, transformer=transformer,
+        )
+
+        # Fix 1: Override _execution_device to force main GPU
+        # Without this, pipeline returns text_encoder's device, causing VAE
+        # to receive tensors on the wrong GPU
+        main_device = torch.device(main_dev)
+        QwenImageEditPipeline._execution_device = property(lambda self: main_device)
+
+        # Fix 2: Monkey-patch prompt encoding to route inputs to text_encoder's device
+        self.pipe._get_qwen_prompt_embeds = types.MethodType(
+            self._patched_get_qwen_prompt_embeds, self.pipe)
+
+        self._loading_strategy = "pinned_multi_gpu"
+        logger.info(f"Pinned 2-GPU pipeline ready")
+        return True
+
+    def load_model(self) -> bool:
+        """Load the model with the best available strategy.
+
+        Strategy priority (GPU strategies always attempted first):
+        1. Pinned 2-GPU: transformer+VAE on large GPU, text_encoder on secondary
+           (requires main GPU >= 42GB, secondary >= 17GB)
+           Benchmark: 169.9s (4.25s/step) - 1.36x
+        2. Balanced single-GPU: device_map="balanced" on single large GPU
+           (requires GPU >= 45GB)
+           Benchmark: 184.4s (4.61s/step) - 1.25x
+        3. CPU offload: model components shuttle between CPU and GPU
+           (requires enable_cpu_offload=True)
+           Benchmark: 231.5s (5.79s/step) - 1.0x baseline
+        4. Direct load: entire model on single GPU (may OOM)
+        """
+        if self._loaded:
+            return True
+
+        try:
+            from diffusers import QwenImageEditPipeline
+
+            model_id = self.MODELS.get(self.model_variant, self.MODELS["full"])
+            main_idx = self._parse_device_idx(self.device)
+            main_vram = self._get_gpu_vram_gb(main_idx)
+            logger.info(f"Loading Qwen-Image-Edit ({self.model_variant}) from {model_id}...")
+            logger.info(f"Main GPU cuda:{main_idx}: {main_vram:.1f} GB VRAM")
+
+            start_time = time.time()
+            loaded = False
+
+            # Strategy 1: Pinned 2-GPU (always try first if main GPU is large enough)
+            if not loaded and main_vram >= self.MAIN_GPU_MIN_VRAM_GB:
+                encoder_idx = None
+                if self.encoder_device:
+                    encoder_idx = self._parse_device_idx(self.encoder_device)
+                    enc_vram = self._get_gpu_vram_gb(encoder_idx)
+                    if enc_vram < self.ENCODER_GPU_MIN_VRAM_GB:
+                        logger.warning(
+                            f"Specified encoder device cuda:{encoder_idx} has "
+                            f"{enc_vram:.1f} GB, need {self.ENCODER_GPU_MIN_VRAM_GB} GB. "
+                            f"Falling back to auto-detect.")
+                        encoder_idx = None
+
+                if encoder_idx is None:
+                    encoder_idx = self._find_encoder_gpu(main_idx)
+
+                if encoder_idx is not None:
+                    self._load_pinned_multi_gpu(model_id, main_idx, encoder_idx)
+                    loaded = True
+
+            # Strategy 2: Balanced single-GPU
+            if not loaded and main_vram >= self.BALANCED_VRAM_THRESHOLD_GB:
+                max_mem_gb = int(main_vram - 4)
+                self.pipe = QwenImageEditPipeline.from_pretrained(
+                    model_id, torch_dtype=self.dtype,
+                    device_map="balanced",
+                    max_memory={main_idx: f"{max_mem_gb}GiB"},
+                )
+                self._loading_strategy = "balanced_single"
+                logger.info(f"Loaded with device_map='balanced', max_memory={max_mem_gb}GiB")
+                loaded = True
+
+            # Strategy 3: CPU offload (only if allowed)
+            if not loaded and self.enable_cpu_offload:
+                self.pipe = QwenImageEditPipeline.from_pretrained(
+                    model_id, torch_dtype=self.dtype)
+                self.pipe.enable_model_cpu_offload()
+                self._loading_strategy = "cpu_offload"
+                logger.info("Loaded with enable_model_cpu_offload()")
+                loaded = True
+
+            # Strategy 4: Direct load (last resort, may OOM)
+            if not loaded:
+                self.pipe = QwenImageEditPipeline.from_pretrained(
+                    model_id, torch_dtype=self.dtype)
+                self.pipe.to(self.device)
+                self._loading_strategy = "direct"
+                logger.info(f"Loaded directly to {self.device}")
+
+            self.pipe.set_progress_bar_config(disable=None)
+
+            load_time = time.time() - start_time
+            logger.info(f"Qwen-Image-Edit loaded in {load_time:.1f}s (strategy: {self._loading_strategy})")
+
+            self._loaded = True
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to load Qwen-Image-Edit: {e}", exc_info=True)
+            return False
+
+    def unload_model(self):
+        """Unload model from memory."""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+            self._loaded = False
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            logger.info("Qwen-Image-Edit-2511 unloaded")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 1.0,
+        true_cfg_scale: float = 4.0
+    ) -> GenerationResult:
+        """
+        Generate/edit image using Qwen-Image-Edit-2511.
+
+        Args:
+            request: GenerationRequest object
+            num_inference_steps: Number of denoising steps
+            guidance_scale: Classifier-free guidance scale
+            true_cfg_scale: True CFG scale for better control
+
+        Returns:
+            GenerationResult object
+        """
+        if not self._loaded:
+            if not self.load_model():
+                return GenerationResult.error_result("Failed to load Qwen-Image-Edit-2511 model")
+
+        try:
+            start_time = time.time()
+
+            # Target dimensions for post-processing crop+resize
+            target_w, target_h = self._get_dimensions(request.aspect_ratio)
+
+            # Build input images list
+            input_images = []
+            if request.has_input_images:
+                input_images = [img for img in request.input_images if img is not None]
+
+            # Always generate at the proven native resolution (1104x1472).
+            # Other resolutions cause VAE tiling artifacts.
+            native_w, native_h = self.NATIVE_RESOLUTION
+            gen_kwargs = {
+                "prompt": request.prompt,
+                "negative_prompt": request.negative_prompt or " ",
+                "height": native_h,
+                "width": native_w,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+                "true_cfg_scale": true_cfg_scale,
+                "num_images_per_prompt": 1,
+                "generator": torch.manual_seed(42),
+            }
+
+            # Qwen-Image-Edit is a single-image editor: use only the first image.
+            # The character service passes multiple references (face, body, costume)
+            # but the costume/view info is already encoded in the text prompt.
+            if input_images:
+                gen_kwargs["image"] = input_images[0]
+
+            logger.info(f"Generating with Qwen-Image-Edit: {request.prompt[:80]}...")
+            logger.info(f"Input images: {len(input_images)} (using first)")
+            logger.info(f"Native: {native_w}x{native_h}, target: {target_w}x{target_h}")
+
+            # Generate at proven native resolution
+            with torch.inference_mode():
+                output = self.pipe(**gen_kwargs)
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Generated in {generation_time:.2f}s: {image.size}")
+
+            # Crop + resize to requested aspect ratio
+            image = self._crop_and_resize(image, target_w, target_h)
+            logger.info(f"Post-processed to: {image.size}")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Generated with Qwen-Image-Edit in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"Qwen-Image-Edit generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Qwen-Image-Edit error: {str(e)}")
+
+    @staticmethod
+    def _crop_and_resize(image: Image.Image, target_w: int, target_h: int) -> Image.Image:
+        """Crop image to target aspect ratio, then resize to target dimensions.
+
+        Centers the crop on the image so equal amounts are trimmed from
+        each side.  Uses LANCZOS for high-quality downscaling.
+        """
+        src_w, src_h = image.size
+        target_ratio = target_w / target_h
+        src_ratio = src_w / src_h
+
+        if abs(target_ratio - src_ratio) < 0.01:
+            # Already the right aspect ratio, just resize
+            return image.resize((target_w, target_h), Image.LANCZOS)
+
+        if target_ratio < src_ratio:
+            # Target is taller/narrower than source → crop sides
+            crop_w = int(src_h * target_ratio)
+            offset = (src_w - crop_w) // 2
+            image = image.crop((offset, 0, offset + crop_w, src_h))
+        else:
+            # Target is wider than source → crop top/bottom
+            crop_h = int(src_w / target_ratio)
+            offset = (src_h - crop_h) // 2
+            image = image.crop((0, offset, src_w, offset + crop_h))
+
+        return image.resize((target_w, target_h), Image.LANCZOS)
+
+    def _get_dimensions(self, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return self.ASPECT_RATIOS.get(ratio, (1024, 1024))
+
+    def is_healthy(self) -> bool:
+        """Check if model is loaded and ready."""
+        return self._loaded and self.pipe is not None
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))

src/utils.py

@@ -0,0 +1,247 @@
+"""
+Utility Functions
+=================
+
+Helper functions for image processing and file operations.
+"""
+
+import re
+import logging
+from pathlib import Path
+from typing import Optional, Union
+from datetime import datetime
+from PIL import Image
+
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_pil_image(
+    obj: Union[Image.Image, str, Path, None],
+    context: str = ""
+) -> Image.Image:
+    """
+    Ensure object is a PIL Image.
+
+    Args:
+        obj: Image, path, or None
+        context: Context for error messages
+
+    Returns:
+        PIL Image
+
+    Raises:
+        ValueError: If object cannot be converted to Image
+    """
+    if obj is None:
+        raise ValueError(f"[{context}] Image is None")
+
+    if isinstance(obj, Image.Image):
+        return obj
+
+    if isinstance(obj, (str, Path)):
+        try:
+            return Image.open(obj)
+        except Exception as e:
+            raise ValueError(f"[{context}] Failed to load image from path: {e}")
+
+    raise ValueError(f"[{context}] Unsupported image type: {type(obj)}")
+
+
+def sanitize_filename(name: str) -> str:
+    """
+    Sanitize string for use as filename.
+
+    Args:
+        name: Original name
+
+    Returns:
+        Safe filename string
+    """
+    # Replace problematic characters
+    safe_name = re.sub(r'[<>:"/\\|?*]', '_', name)
+    # Remove leading/trailing spaces and dots
+    safe_name = safe_name.strip('. ')
+    # Limit length
+    if len(safe_name) > 100:
+        safe_name = safe_name[:100]
+    return safe_name or "unnamed"
+
+
+def save_image(
+    image: Image.Image,
+    directory: Path,
+    base_name: str,
+    format: str = "PNG"
+) -> Path:
+    """
+    Save image to directory.
+
+    Args:
+        image: PIL Image to save
+        directory: Output directory
+        base_name: Base filename (without extension)
+        format: Image format
+
+    Returns:
+        Path to saved file
+    """
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    safe_name = sanitize_filename(base_name)
+    ext = format.lower()
+
+    filename = f"{safe_name}_{timestamp}.{ext}"
+    filepath = directory / filename
+
+    image.save(filepath, format=format)
+    logger.info(f"Saved: {filepath}")
+
+    return filepath
+
+
+def resize_for_display(
+    image: Image.Image,
+    max_size: int = 1024
+) -> Image.Image:
+    """
+    Resize image for display while maintaining aspect ratio.
+
+    Args:
+        image: PIL Image
+        max_size: Maximum dimension
+
+    Returns:
+        Resized image
+    """
+    width, height = image.size
+
+    if width <= max_size and height <= max_size:
+        return image
+
+    if width > height:
+        new_width = max_size
+        new_height = int(height * max_size / width)
+    else:
+        new_height = max_size
+        new_width = int(width * max_size / height)
+
+    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+
+def get_image_info(image: Image.Image) -> str:
+    """Get human-readable image info string."""
+    return f"{image.size[0]}x{image.size[1]} {image.mode}"
+
+
+def preprocess_input_image(
+    image: Image.Image,
+    max_size: int = 1024,
+    target_size: tuple = None,
+    ensure_rgb: bool = True
+) -> Image.Image:
+    """
+    Preprocess input image for model consumption.
+
+    Handles various formats (JFIF, TIFF, WebP, etc.) by converting to RGB PNG-compatible format.
+
+    Args:
+        image: PIL Image to preprocess
+        max_size: Maximum dimension (used if target_size not specified)
+        target_size: Specific (width, height) to resize to
+        ensure_rgb: Convert to RGB mode
+
+    Returns:
+        Preprocessed PIL Image in RGB format
+    """
+    # Ensure we have a copy to avoid modifying original
+    img = image.copy()
+
+    # Force re-encode as PNG-compatible by saving to memory and reloading
+    # This handles weird formats like JFIF, TIFF, etc.
+    import io
+    buf = io.BytesIO()
+
+    # Convert to RGB first if needed
+    if img.mode not in ('RGB', 'RGBA'):
+        img = img.convert('RGB')
+
+    # Save as PNG to buffer and reload - this normalizes the format
+    img.save(buf, format='PNG')
+    buf.seek(0)
+    img = Image.open(buf)
+    img.load()  # Force load into memory
+
+    # Convert to RGB if needed (handle RGBA)
+    if ensure_rgb and img.mode != 'RGB':
+        if img.mode == 'RGBA':
+            # Handle transparency by compositing on white background
+            background = Image.new('RGB', img.size, (255, 255, 255))
+            background.paste(img, mask=img.split()[3])
+            img = background
+        else:
+            img = img.convert('RGB')
+
+    # Resize to target size or max_size
+    if target_size:
+        img = img.resize(target_size, Image.Resampling.LANCZOS)
+    else:
+        width, height = img.size
+        if width > max_size or height > max_size:
+            if width > height:
+                new_width = max_size
+                new_height = int(height * max_size / width)
+            else:
+                new_height = max_size
+                new_width = int(width * max_size / height)
+            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    return img
+
+
+def preprocess_images_for_backend(
+    images: list,
+    backend_type: str,
+    aspect_ratio: str = "1:1"
+) -> list:
+    """
+    Preprocess a list of images for a specific backend.
+
+    Args:
+        images: List of PIL Images
+        backend_type: Backend type string (e.g., 'flux_klein', 'qwen_comfyui')
+        aspect_ratio: Target aspect ratio
+
+    Returns:
+        List of preprocessed PIL Images
+    """
+    if not images:
+        return images
+
+    # Backend-specific settings
+    # FLUX models work best with smaller input images (512-768px)
+    backend_configs = {
+        'flux_klein': {'max_size': 768},           # 4B - faster with smaller inputs
+        'flux_klein_9b_fp8': {'max_size': 768},    # 9B - same, quality comes from model not input size
+        'qwen_image_edit': {'max_size': 1024},
+        'qwen_comfyui': {'max_size': 1024},
+        'zimage_turbo': {'max_size': 768},
+        'zimage_base': {'max_size': 768},
+        'longcat_edit': {'max_size': 768},
+        'gemini_flash': {'max_size': 1024},        # Gemini handles larger but 1024 is fine
+        'gemini_pro': {'max_size': 1024},
+    }
+
+    config = backend_configs.get(backend_type, {'max_size': 1024})
+    max_size = config['max_size']
+
+    processed = []
+    for img in images:
+        if img is not None:
+            processed.append(preprocess_input_image(img, max_size=max_size))
+        else:
+            processed.append(None)
+
+    return processed

src/zimage_client.py

@@ -0,0 +1,351 @@
+"""
+Z-Image Client
+==============
+
+Client for Z-Image (Tongyi-MAI) local image generation.
+Supports text-to-image and image-to-image editing.
+
+Z-Image is a 6B parameter model that achieves state-of-the-art quality
+with only 8-9 inference steps, fitting in 16GB VRAM.
+"""
+
+import logging
+import time
+from typing import Optional, List
+from PIL import Image
+
+import torch
+
+from .models import GenerationRequest, GenerationResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class ZImageClient:
+    """
+    Client for Z-Image models from Tongyi-MAI.
+
+    Supports:
+    - Text-to-image generation (ZImagePipeline)
+    - Image-to-image editing (ZImageImg2ImgPipeline)
+    - Multiple model variants (Turbo, Base, Edit, Omni)
+    """
+
+    # Model variants
+    MODELS = {
+        # Turbo - Fast, distilled, 8-9 steps, fits 16GB VRAM
+        "turbo": "Tongyi-MAI/Z-Image-Turbo",
+        # Base - Quality-focused, more steps
+        "base": "Tongyi-MAI/Z-Image",
+        # Edit - Fine-tuned for instruction-following image editing
+        "edit": "Tongyi-MAI/Z-Image-Edit",
+        # Omni - Versatile, supports both generation and editing
+        "omni": "Tongyi-MAI/Z-Image-Omni-Base",
+    }
+
+    # Aspect ratio to dimensions mapping
+    # Z-Image supports 512x512 to 2048x2048
+    ASPECT_RATIOS = {
+        "1:1": (1024, 1024),
+        "16:9": (1344, 768),
+        "9:16": (768, 1344),
+        "21:9": (1536, 640),    # Cinematic ultra-wide
+        "3:2": (1248, 832),
+        "2:3": (832, 1248),
+        "3:4": (896, 1152),
+        "4:3": (1152, 896),
+        "4:5": (896, 1120),
+        "5:4": (1120, 896),
+    }
+
+    # Default settings for each model variant
+    MODEL_DEFAULTS = {
+        "turbo": {"steps": 9, "guidance": 0.0},      # Fast, no CFG needed
+        "base": {"steps": 50, "guidance": 4.0},      # Quality-focused
+        "edit": {"steps": 28, "guidance": 3.5},      # Editing
+        "omni": {"steps": 28, "guidance": 3.5},      # Versatile
+    }
+
+    def __init__(
+        self,
+        model_variant: str = "turbo",
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        enable_cpu_offload: bool = True,
+    ):
+        """
+        Initialize Z-Image client.
+
+        Args:
+            model_variant: Model variant to use:
+                - "turbo": Fast, 9 steps, 16GB VRAM (RECOMMENDED)
+                - "base": Quality-focused, 50 steps
+                - "edit": Instruction-following image editing
+                - "omni": Versatile generation + editing
+            device: Device to use (cuda or cpu)
+            dtype: Data type for model weights (bfloat16 recommended)
+            enable_cpu_offload: Enable CPU offload to save VRAM
+        """
+        self.model_variant = model_variant
+        self.device = device
+        self.dtype = dtype
+        self.enable_cpu_offload = enable_cpu_offload
+        self.pipe = None
+        self.pipe_img2img = None
+        self._loaded = False
+
+        # Get default settings for this variant
+        defaults = self.MODEL_DEFAULTS.get(model_variant, {"steps": 9, "guidance": 0.0})
+        self.default_steps = defaults["steps"]
+        self.default_guidance = defaults["guidance"]
+
+        logger.info(f"ZImageClient initialized (variant: {model_variant}, steps: {self.default_steps}, guidance: {self.default_guidance})")
+
+    def load_model(self) -> bool:
+        """Load the model into memory."""
+        if self._loaded:
+            return True
+
+        try:
+            # Get model ID for selected variant
+            model_id = self.MODELS.get(self.model_variant, self.MODELS["turbo"])
+            logger.info(f"Loading Z-Image ({self.model_variant}) from {model_id}...")
+
+            start_time = time.time()
+
+            # Import diffusers pipelines for Z-Image
+            # Requires latest diffusers: pip install git+https://github.com/huggingface/diffusers
+            from diffusers import ZImagePipeline, ZImageImg2ImgPipeline
+
+            # Load text-to-image pipeline
+            self.pipe = ZImagePipeline.from_pretrained(
+                model_id,
+                torch_dtype=self.dtype,
+            )
+
+            # Load img2img pipeline (shares components)
+            self.pipe_img2img = ZImageImg2ImgPipeline.from_pretrained(
+                model_id,
+                torch_dtype=self.dtype,
+                # Share components to save memory
+                text_encoder=self.pipe.text_encoder,
+                tokenizer=self.pipe.tokenizer,
+                vae=self.pipe.vae,
+                transformer=self.pipe.transformer,
+                scheduler=self.pipe.scheduler,
+            )
+
+            # Apply memory optimization
+            if self.enable_cpu_offload:
+                self.pipe.enable_model_cpu_offload()
+                self.pipe_img2img.enable_model_cpu_offload()
+                logger.info("CPU offload enabled")
+            else:
+                self.pipe.to(self.device)
+                self.pipe_img2img.to(self.device)
+                logger.info(f"Model moved to {self.device}")
+
+            # Optional: Enable flash attention if available
+            try:
+                self.pipe.transformer.set_attention_backend("flash")
+                self.pipe_img2img.transformer.set_attention_backend("flash")
+                logger.info("Flash Attention enabled")
+            except Exception:
+                logger.info("Flash Attention not available, using default SDPA")
+
+            load_time = time.time() - start_time
+            logger.info(f"Z-Image ({self.model_variant}) loaded in {load_time:.1f}s")
+
+            # Validate by running a test generation
+            logger.info("Validating model with test generation...")
+            try:
+                test_result = self.pipe(
+                    prompt="A simple test image",
+                    height=256,
+                    width=256,
+                    guidance_scale=0.0,
+                    num_inference_steps=2,
+                    generator=torch.Generator(device="cpu").manual_seed(42),
+                )
+                if test_result.images[0] is not None:
+                    logger.info("Model validation successful")
+                else:
+                    logger.error("Model validation failed: no output image")
+                    return False
+            except Exception as e:
+                logger.error(f"Model validation failed: {e}", exc_info=True)
+                return False
+
+            self._loaded = True
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to load Z-Image: {e}", exc_info=True)
+            return False
+
+    def unload_model(self):
+        """Unload model from memory."""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+        if self.pipe_img2img is not None:
+            del self.pipe_img2img
+            self.pipe_img2img = None
+
+        self._loaded = False
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        logger.info("Z-Image unloaded")
+
+    def generate(
+        self,
+        request: GenerationRequest,
+        num_inference_steps: int = None,
+        guidance_scale: float = None
+    ) -> GenerationResult:
+        """
+        Generate image using Z-Image.
+
+        Args:
+            request: GenerationRequest object
+            num_inference_steps: Number of denoising steps (9 for turbo)
+            guidance_scale: Classifier-free guidance scale (0.0 for turbo)
+
+        Returns:
+            GenerationResult object
+        """
+        if not self._loaded:
+            if not self.load_model():
+                return GenerationResult.error_result("Failed to load Z-Image model")
+
+        # Use model defaults if not specified
+        if num_inference_steps is None:
+            num_inference_steps = self.default_steps
+        if guidance_scale is None:
+            guidance_scale = self.default_guidance
+
+        try:
+            start_time = time.time()
+
+            # Get dimensions from aspect ratio
+            width, height = self._get_dimensions(request.aspect_ratio)
+
+            logger.info(f"Generating with Z-Image {self.model_variant}: steps={num_inference_steps}, guidance={guidance_scale}")
+
+            # Check if we have input images (use img2img pipeline)
+            if request.has_input_images:
+                return self._generate_img2img(
+                    request, width, height, num_inference_steps, guidance_scale, start_time
+                )
+
+            # Text-to-image generation
+            gen_kwargs = {
+                "prompt": request.prompt,
+                "height": height,
+                "width": width,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps,
+                "generator": torch.Generator(device="cpu").manual_seed(42),
+            }
+
+            # Add negative prompt if present
+            if request.negative_prompt:
+                gen_kwargs["negative_prompt"] = request.negative_prompt
+
+            logger.info(f"Generating with Z-Image: {request.prompt[:80]}...")
+
+            # Generate
+            with torch.inference_mode():
+                output = self.pipe(**gen_kwargs)
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Generated in {generation_time:.2f}s: {image.size}")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Generated with Z-Image ({self.model_variant}) in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"Z-Image generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Z-Image error: {str(e)}")
+
+    def _generate_img2img(
+        self,
+        request: GenerationRequest,
+        width: int,
+        height: int,
+        num_inference_steps: int,
+        guidance_scale: float,
+        start_time: float
+    ) -> GenerationResult:
+        """Generate using img2img pipeline with input images."""
+        try:
+            # Get the first valid input image
+            input_image = None
+            for img in request.input_images:
+                if img is not None:
+                    input_image = img
+                    break
+
+            if input_image is None:
+                return GenerationResult.error_result("No valid input image provided")
+
+            # Resize input image to target dimensions
+            input_image = input_image.resize((width, height), Image.Resampling.LANCZOS)
+
+            # Build generation kwargs for img2img
+            gen_kwargs = {
+                "prompt": request.prompt,
+                "image": input_image,
+                "strength": 0.6,  # How much to transform the image
+                "height": height,
+                "width": width,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps,
+                "generator": torch.Generator(device="cpu").manual_seed(42),
+            }
+
+            # Add negative prompt if present
+            if request.negative_prompt:
+                gen_kwargs["negative_prompt"] = request.negative_prompt
+
+            logger.info(f"Generating img2img with Z-Image: {request.prompt[:80]}...")
+
+            # Generate
+            with torch.inference_mode():
+                output = self.pipe_img2img(**gen_kwargs)
+                image = output.images[0]
+
+            generation_time = time.time() - start_time
+            logger.info(f"Generated img2img in {generation_time:.2f}s: {image.size}")
+
+            return GenerationResult.success_result(
+                image=image,
+                message=f"Generated with Z-Image img2img ({self.model_variant}) in {generation_time:.2f}s",
+                generation_time=generation_time
+            )
+
+        except Exception as e:
+            logger.error(f"Z-Image img2img generation failed: {e}", exc_info=True)
+            return GenerationResult.error_result(f"Z-Image img2img error: {str(e)}")
+
+    def _get_dimensions(self, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return self.ASPECT_RATIOS.get(ratio, (1024, 1024))
+
+    def is_healthy(self) -> bool:
+        """Check if model is loaded and ready."""
+        return self._loaded and self.pipe is not None
+
+    @classmethod
+    def get_dimensions(cls, aspect_ratio: str) -> tuple:
+        """Get pixel dimensions for aspect ratio."""
+        ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio
+        return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))