yamlgraph 0.1.1__py3-none-any.whl

examples/storyboard/nodes/animated_image_node.py

@@ -0,0 +1,138 @@
+"""Animated storyboard node for generating frame images.
+
+Generates 3 images per panel: first_frame, original, last_frame.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .replicate_tool import generate_image
+
+logger = logging.getLogger(__name__)
+
+GraphState = dict[str, Any]
+
+
+def generate_animated_images_node(state: GraphState) -> dict:
+    """Generate images for all animated panel frames.
+
+    Reads animated_panels from state and generates 3 images per panel.
+    Saves to outputs/storyboard/{thread_id}/animated/ directory.
+
+    Args:
+        state: Graph state with 'animated_panels' list of {first_frame, original, last_frame}
+
+    Returns:
+        State update with 'images' list organized by panel
+    """
+    animated_panels = state.get("animated_panels", [])
+    if not animated_panels:
+        logger.error("No animated_panels in state")
+        return {
+            "current_step": "generate_animated_images",
+            "images": [],
+            "error": "No animated panels to generate",
+        }
+
+    # Sort by _map_index if present to maintain order
+    if animated_panels and isinstance(animated_panels[0], dict):
+        animated_panels = sorted(
+            animated_panels,
+            key=lambda x: x.get("_map_index", 0) if isinstance(x, dict) else 0,
+        )
+
+    # Create output directory
+    thread_id = state.get("thread_id", datetime.now().strftime("%Y%m%d_%H%M%S"))
+    output_dir = Path("outputs/storyboard") / thread_id / "animated"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    total_images = len(animated_panels) * 3
+    logger.info(
+        f"🎬 Generating {total_images} images ({len(animated_panels)} panels × 3 frames)"
+    )
+
+    # Get model selection from state (default: z-image)
+    model_name = state.get("model", "z-image")
+    logger.info(f"🖼️  Using model: {model_name}")
+
+    # Generate images for each panel
+    all_results: list[dict] = []
+    frame_keys = ["first_frame", "original", "last_frame"]
+
+    for panel_idx, panel in enumerate(animated_panels, 1):
+        # Handle Pydantic model or dict
+        if hasattr(panel, "model_dump"):
+            panel_dict = panel.model_dump()
+        elif isinstance(panel, dict):
+            panel_dict = panel
+        else:
+            logger.warning(f"Panel {panel_idx} has unexpected type: {type(panel)}")
+            continue
+
+        panel_result = {"panel": panel_idx, "frames": {}}
+
+        for frame_key in frame_keys:
+            prompt = panel_dict.get(frame_key, "")
+            if not prompt:
+                logger.warning(f"Panel {panel_idx} missing {frame_key}")
+                continue
+
+            output_path = output_dir / f"panel_{panel_idx}_{frame_key}.png"
+            logger.info(f"📸 Panel {panel_idx} {frame_key}: {prompt[:50]}...")
+
+            result = generate_image(prompt, output_path, model_name=model_name)
+
+            if result.success and result.path:
+                panel_result["frames"][frame_key] = result.path
+            else:
+                logger.error(f"Panel {panel_idx} {frame_key} failed: {result.error}")
+                panel_result["frames"][frame_key] = None
+
+        all_results.append(panel_result)
+
+    # Save metadata
+    story = state.get("story", {})
+    if hasattr(story, "model_dump"):
+        story_dict = story.model_dump()
+    elif isinstance(story, dict):
+        story_dict = story
+    else:
+        story_dict = {}
+
+    metadata_path = output_dir / "animated_story.json"
+    metadata = {
+        "concept": state.get("concept", ""),
+        "title": story_dict.get("title", ""),
+        "narrative": story_dict.get("narrative", ""),
+        "panels": [
+            {
+                "index": r["panel"],
+                "frames": r["frames"],
+                "prompts": {
+                    k: animated_panels[r["panel"] - 1].get(k, "")
+                    if isinstance(animated_panels[r["panel"] - 1], dict)
+                    else ""
+                    for k in frame_keys
+                },
+            }
+            for r in all_results
+        ],
+        "generated_at": datetime.now().isoformat(),
+    }
+    metadata_path.write_text(json.dumps(metadata, indent=2))
+    logger.info(f"📝 Metadata saved: {metadata_path}")
+
+    # Count successes
+    success_count = sum(1 for r in all_results for path in r["frames"].values() if path)
+    logger.info(f"✅ Generated {success_count}/{total_images} images")
+
+    return {
+        "current_step": "generate_animated_images",
+        "images": all_results,
+        "output_dir": str(output_dir),
+    }

examples/storyboard/nodes/character_node.py

@@ -0,0 +1,162 @@
+"""Character-consistent storyboard node.
+
+This node:
+1. Generates a character image from description (step 0)
+2. Uses image-to-image editing to place character in each panel scene
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .replicate_tool import ImageResult, edit_image, generate_image
+
+logger = logging.getLogger(__name__)
+
+# Type alias for state
+GraphState = dict[str, Any]
+
+
+def generate_character_storyboard(state: GraphState) -> dict:
+    """Generate character-consistent storyboard images.
+
+    Step 0: Generate base character image from character_prompt
+    Panels 1-3: Use image-to-image to place character in each scene
+
+    Args:
+        state: Graph state with 'story' containing character and panel prompts
+
+    Returns:
+        State update with 'images' list and metadata
+    """
+    story = state.get("story")
+    if not story:
+        logger.error("No story in state")
+        return {
+            "current_step": "generate_character_storyboard",
+            "images": [],
+            "error": "No story in state",
+        }
+
+    # Handle Pydantic model or dict
+    if hasattr(story, "model_dump"):
+        story_dict = story.model_dump()
+    elif isinstance(story, dict):
+        story_dict = story
+    else:
+        story_dict = {}
+
+    # Extract prompts
+    character_prompt = story_dict.get("character_prompt", "")
+    panels = story_dict.get("panels", [])
+
+    if not character_prompt:
+        logger.error("No character_prompt in story")
+        return {
+            "current_step": "generate_character_storyboard",
+            "images": [],
+            "error": "No character_prompt provided",
+        }
+
+    if not panels:
+        logger.error("No panels in story")
+        return {
+            "current_step": "generate_character_storyboard",
+            "images": [],
+            "error": "No panel prompts provided",
+        }
+
+    # Create output directory
+    thread_id = state.get("thread_id", datetime.now().strftime("%Y%m%d_%H%M%S"))
+    output_dir = Path("outputs/storyboard") / thread_id
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    logger.info(f"🎬 Generating character-consistent storyboard in {output_dir}")
+
+    # Get model selection from state (default: z-image for character)
+    model_name = state.get("model", "z-image")
+    logger.info(f"🖼️  Using model for character: {model_name}")
+
+    image_paths: list[str] = []
+    results: list[ImageResult] = []
+
+    # Step 0: Generate base character image
+    character_path = output_dir / "character.png"
+    logger.info(f"👤 Step 0 - Creating character: {character_prompt[:60]}...")
+
+    character_result = generate_image(
+        prompt=character_prompt,
+        output_path=character_path,
+        model_name=model_name,
+    )
+    results.append(character_result)
+
+    if not character_result.success:
+        logger.error(f"Character generation failed: {character_result.error}")
+        return {
+            "current_step": "generate_character_storyboard",
+            "images": [],
+            "error": f"Character generation failed: {character_result.error}",
+        }
+
+    image_paths.append(str(character_path))
+    logger.info(f"✓ Character created: {character_path}")
+
+    # Panels 1-3: Image-to-image editing with character as base
+    for i, panel_prompt in enumerate(panels[:3], 1):  # Max 3 panels
+        if not panel_prompt:
+            logger.warning(f"Panel {i} has no prompt, skipping")
+            continue
+
+        panel_path = output_dir / f"panel_{i}.png"
+        logger.info(f"📸 Panel {i}: {panel_prompt[:60]}...")
+
+        panel_result = edit_image(
+            input_image=character_path,
+            prompt=panel_prompt,
+            output_path=panel_path,
+            aspect_ratio="16:9",
+        )
+        results.append(panel_result)
+
+        if panel_result.success and panel_result.path:
+            image_paths.append(panel_result.path)
+            logger.info(f"✓ Panel {i} created")
+        else:
+            logger.error(f"Panel {i} failed: {panel_result.error}")
+
+    # Save metadata
+    metadata_path = output_dir / "story.json"
+    metadata = {
+        "concept": state.get("concept", ""),
+        "title": story_dict.get("title", ""),
+        "narrative": story_dict.get("narrative", ""),
+        "character_prompt": character_prompt,
+        "character_image": str(character_path),
+        "panels": [
+            {
+                "prompt": panels[i] if i < len(panels) else "",
+                "image": image_paths[i + 1] if i + 1 < len(image_paths) else None,
+            }
+            for i in range(len(panels[:3]))
+        ],
+        "generated_at": datetime.now().isoformat(),
+    }
+    metadata_path.write_text(json.dumps(metadata, indent=2))
+    logger.info(f"📝 Metadata saved: {metadata_path}")
+
+    success_count = sum(1 for r in results if r.success)
+    logger.info(
+        f"✅ Generated {success_count}/{len(results)} images (1 character + {len(panels[:3])} panels)"
+    )
+
+    return {
+        "current_step": "generate_character_storyboard",
+        "images": image_paths,
+        "character_image": str(character_path),
+        "output_dir": str(output_dir),
+    }

examples/storyboard/nodes/image_node.py

@@ -0,0 +1,118 @@
+"""Storyboard node for generating panel images.
+
+This node takes story panels from the LLM and generates images via Replicate.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .replicate_tool import ImageResult, generate_image
+
+logger = logging.getLogger(__name__)
+
+# Type alias for state
+GraphState = dict[str, Any]
+
+
+def generate_images_node(state: GraphState) -> dict:
+    """Generate images for each story panel.
+
+    Reads panel prompts from state.story and generates images.
+    Saves to outputs/storyboard/{thread_id}/ directory.
+
+    Args:
+        state: Graph state with 'story' containing panel prompts
+
+    Returns:
+        State update with 'images' list and metadata
+    """
+    story = state.get("story")
+    if not story:
+        logger.error("No story in state")
+        return {
+            "current_step": "generate_images",
+            "images": [],
+            "error": "No story panels to generate",
+        }
+
+    # Handle Pydantic model or dict
+    if hasattr(story, "model_dump"):
+        story_dict = story.model_dump()
+    elif isinstance(story, dict):
+        story_dict = story
+    else:
+        story_dict = {"panels": [str(story)]}
+
+    # Extract panel prompts (supports dynamic list)
+    panels = story_dict.get("panels", [])
+    if not panels:
+        # Fallback for legacy panel_1/2/3 format
+        panels = [
+            story_dict.get("panel_1", ""),
+            story_dict.get("panel_2", ""),
+            story_dict.get("panel_3", ""),
+        ]
+        panels = [p for p in panels if p]  # Remove empty
+
+    # Create output directory
+    thread_id = state.get("thread_id", datetime.now().strftime("%Y%m%d_%H%M%S"))
+    output_dir = Path("outputs/storyboard") / thread_id
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    logger.info(f"🎬 Generating {len(panels)}-panel storyboard in {output_dir}")
+
+    # Get model selection from state (default: z-image)
+    model_name = state.get("model", "z-image")
+    logger.info(f"\ud83d\uddbc\ufe0f  Using model: {model_name}")
+
+    # Generate each panel image
+    results: list[ImageResult] = []
+    image_paths: list[str] = []
+
+    for i, prompt in enumerate(panels, 1):
+        if not prompt:
+            logger.warning(f"Panel {i} has no prompt, skipping")
+            continue
+
+        output_path = output_dir / f"panel_{i}.png"
+        logger.info(f"\ud83d\udcf8 Panel {i}: {prompt[:60]}...")
+
+        result = generate_image(prompt, output_path, model_name=model_name)
+        results.append(result)
+
+        if result.success and result.path:
+            image_paths.append(result.path)
+        else:
+            logger.error(f"Panel {i} failed: {result.error}")
+
+    # Save story metadata
+    metadata_path = output_dir / "story.json"
+    metadata = {
+        "concept": state.get("concept", ""),
+        "title": story_dict.get("title", ""),
+        "narrative": story_dict.get("narrative", ""),
+        "panels": [
+            {
+                "prompt": panels[i] if i < len(panels) else "",
+                "image": image_paths[i] if i < len(image_paths) else None,
+            }
+            for i in range(max(len(panels), len(image_paths)))
+        ],
+        "generated_at": datetime.now().isoformat(),
+    }
+    metadata_path.write_text(json.dumps(metadata, indent=2))
+    logger.info(f"📝 Metadata saved: {metadata_path}")
+
+    success_count = sum(1 for r in results if r.success)
+    logger.info(f"✅ Generated {success_count}/{len(panels)} images")
+
+    return {
+        "current_step": "generate_images",
+        "images": image_paths,
+        "output_dir": str(output_dir),
+    }

examples/storyboard/nodes/replicate_tool.py

@@ -0,0 +1,238 @@
+"""Replicate image generation tool for storyboard workflow.
+
+Supports multiple models:
+- z-image: Fast, good for realistic/photographic (default)
+- hidream: Better for cartoons, illustrations, stylized art
+- p-image-edit: Image-to-image editing for character consistency
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Model configurations
+MODELS = {
+    "z-image": {
+        "id": "prunaai/z-image-turbo",
+        "width": 1344,
+        "height": 768,
+        "params": {
+            "guidance_scale": 0,
+            "num_inference_steps": 8,
+        },
+    },
+    "hidream": {
+        "id": "prunaai/hidream-l1-fast:f67f0ec7ef9fe91b74e8a68d34efaa9145bec28675cb190cbff8a70f0490256e",
+        "resolution": "1360 \u00d7 768 (Landscape)",
+        "params": {
+            "model_type": "fast",
+            "speed_mode": "Juiced \U0001f525 (more speed)",
+        },
+    },
+}
+
+DEFAULT_MODEL = "z-image"
+
+# Check if replicate is available
+try:
+    import replicate
+
+    REPLICATE_AVAILABLE = True
+except ImportError:
+    REPLICATE_AVAILABLE = False
+    logger.warning("replicate package not installed. Run: pip install replicate")
+
+
+@dataclass
+class ImageResult:
+    """Result from image generation."""
+
+    success: bool
+    path: str | None = None
+    error: str | None = None
+
+
+def generate_image(
+    prompt: str,
+    output_path: str | Path,
+    model_name: str = DEFAULT_MODEL,
+) -> ImageResult:
+    """Generate an image using Replicate API.
+
+    Args:
+        prompt: Text prompt for image generation
+        output_path: Path to save the generated image
+        model_name: Model to use ('z-image' or 'hidream')
+
+    Returns:
+        ImageResult with success status and path or error
+    """
+    if not REPLICATE_AVAILABLE:
+        return ImageResult(success=False, error="replicate package not installed")
+
+    api_token = os.environ.get("REPLICATE_API_TOKEN")
+    if not api_token:
+        return ImageResult(
+            success=False, error="REPLICATE_API_TOKEN not set in environment"
+        )
+
+    # Get model config
+    model_config = MODELS.get(model_name, MODELS[DEFAULT_MODEL])
+    model_id = model_config["id"]
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        logger.info(f"🎨 Generating image with {model_name}: {prompt[:50]}...")
+
+        # Build input params based on model
+        if model_name == "hidream":
+            input_params = {
+                "prompt": prompt,
+                "seed": -1,
+                "resolution": model_config["resolution"],
+                "output_format": "png",
+                "output_quality": 80,
+                "disable_safety_checker": True,
+                **model_config["params"],
+            }
+        else:
+            # z-image and default
+            input_params = {
+                "prompt": prompt,
+                "width": model_config.get("width", 1344),
+                "height": model_config.get("height", 768),
+                "output_format": "png",
+                "output_quality": 80,
+                "disable_safety_checker": True,
+                **model_config.get("params", {}),
+            }
+
+        # Run the model
+        client = replicate.Client(api_token=api_token)
+        output = client.run(model_id, input=input_params)
+
+        # output is typically a URL or file-like object
+        image_url = output if isinstance(output, str) else str(output)
+
+        # Download the image
+        logger.info(f"📥 Downloading image to {output_path}")
+        response = httpx.get(image_url, timeout=60.0)
+        response.raise_for_status()
+
+        output_path.write_bytes(response.content)
+        logger.info(f"✓ Image saved: {output_path}")
+
+        return ImageResult(success=True, path=str(output_path))
+
+    except Exception as e:
+        logger.error(f"Image generation failed: {e}")
+        return ImageResult(success=False, error=str(e))
+
+
+def edit_image(
+    input_image: str | Path,
+    prompt: str,
+    output_path: str | Path,
+    aspect_ratio: str = "16:9",
+    turbo: bool = True,
+    magic: float | None = None,
+) -> ImageResult:
+    """Edit an image using Replicate p-image-edit model.
+
+    Uses the input image as base and applies the prompt as modifications.
+    Great for maintaining character consistency across panels.
+
+    Args:
+        input_image: Path to the source image
+        prompt: Edit instructions (what to change/add)
+        output_path: Path to save the edited image
+        aspect_ratio: Output aspect ratio (default 16:9)
+        turbo: Use turbo mode for faster generation
+        magic: Prompt strength 0-1 (lower = more original, higher = more prompt)
+
+    Returns:
+        ImageResult with success status and path or error
+    """
+    if not REPLICATE_AVAILABLE:
+        return ImageResult(success=False, error="replicate package not installed")
+
+    api_token = os.environ.get("REPLICATE_API_TOKEN")
+    if not api_token:
+        return ImageResult(
+            success=False, error="REPLICATE_API_TOKEN not set in environment"
+        )
+
+    input_image = Path(input_image)
+    if not input_image.exists():
+        return ImageResult(success=False, error=f"Input image not found: {input_image}")
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        logger.info(f"✏️ Editing image: {prompt[:50]}...")
+
+        client = replicate.Client(api_token=api_token)
+
+        with open(input_image, "rb") as f:
+            input_params = {
+                "turbo": turbo,
+                "images": [f],
+                "prompt": prompt,
+                "aspect_ratio": aspect_ratio,
+                "disable_safety_checker": True,
+            }
+            if magic is not None:
+                input_params["magic"] = magic
+
+            output = client.run(
+                "prunaai/p-image-edit",
+                input=input_params,
+            )
+
+        # Save the output
+        with open(output_path, "wb") as out:
+            out.write(output.read())
+
+        logger.info(f"✓ Edited image saved: {output_path}")
+        return ImageResult(success=True, path=str(output_path))
+
+    except Exception as e:
+        logger.error(f"Image editing failed: {e}")
+        return ImageResult(success=False, error=str(e))
+
+
+def generate_storyboard_images(
+    panel_prompts: list[str],
+    output_dir: str | Path,
+    prefix: str = "panel",
+) -> list[ImageResult]:
+    """Generate multiple images for a storyboard.
+
+    Args:
+        panel_prompts: List of prompts for each panel
+        output_dir: Directory to save images
+        prefix: Filename prefix
+
+    Returns:
+        List of ImageResult for each panel
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    results = []
+    for i, prompt in enumerate(panel_prompts, 1):
+        output_path = output_dir / f"{prefix}_{i}.png"
+        result = generate_image(prompt, output_path)
+        results.append(result)
+
+    return results