yamlgraph 0.3.9__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,49 @@
+"""Storyboard-specific wrappers for image generation.
+
+Re-exports from examples.shared.replicate_tool and adds storyboard-specific functions.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from examples.shared.replicate_tool import (
+    ImageResult,
+    edit_image,
+    generate_image,
+)
+
+# Re-export for backward compatibility
+__all__ = [
+    "ImageResult",
+    "generate_image",
+    "edit_image",
+    "generate_storyboard_images",
+]
+
+
+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

examples/storyboard/retry_images.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Retry image generation from existing animated storyboard metadata.
+
+Usage:
+    python examples/storyboard/retry_images.py outputs/storyboard/20260117_112419/animated
+
+Options:
+    --model MODEL         Image model to use (default: hidream)
+    --reference PATH      Override reference image path
+    --magic FLOAT         Magic value for img2img (default: 0.25)
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+# Add project root to path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+from examples.storyboard.nodes.animated_character_node import (
+    generate_animated_character_images,
+)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Retry image generation from existing metadata"
+    )
+    parser.add_argument(
+        "output_dir",
+        type=Path,
+        help="Path to animated output directory (contains animated_character_story.json)",
+    )
+    parser.add_argument(
+        "--model",
+        default="hidream",
+        help="Image model to use (default: hidream)",
+    )
+    parser.add_argument(
+        "--reference",
+        type=Path,
+        help="Override reference image path",
+    )
+    parser.add_argument(
+        "--new-id",
+        help="New thread ID for output (default: adds _retry suffix)",
+    )
+    args = parser.parse_args()
+
+    # Find metadata file
+    output_dir = args.output_dir
+    if output_dir.name != "animated":
+        output_dir = output_dir / "animated"
+
+    metadata_path = output_dir / "animated_character_story.json"
+    if not metadata_path.exists():
+        print(f"❌ Metadata not found: {metadata_path}")
+        sys.exit(1)
+
+    print(f"📂 Loading: {metadata_path}")
+    metadata = json.loads(metadata_path.read_text())
+
+    # Reconstruct animated_panels from metadata
+    animated_panels = []
+    for panel in metadata["panels"]:
+        animated_panels.append(
+            {
+                "_map_index": panel["index"] - 1,
+                **panel["prompts"],
+            }
+        )
+
+    # Determine thread_id
+    original_thread = output_dir.parent.name
+    thread_id = args.new_id or f"{original_thread}_retry"
+
+    # Build state
+    state = {
+        "concept": metadata["concept"],
+        "model": args.model,
+        "thread_id": thread_id,
+        "story": {
+            "title": metadata["title"],
+            "narrative": metadata["narrative"],
+            "character_prompt": metadata["character_prompt"],
+        },
+        "animated_panels": animated_panels,
+    }
+
+    # Handle reference image
+    if args.reference:
+        state["reference_image"] = str(args.reference)
+    elif metadata.get("reference_image"):
+        ref_path = Path(metadata["reference_image"])
+        if ref_path.exists():
+            state["reference_image"] = str(ref_path)
+            print(f"🎭 Using existing reference: {ref_path}")
+
+    print("🎬 Retrying image generation...")
+    print(f"   Model: {args.model}")
+    print(f"   Panels: {len(animated_panels)}")
+    print(f"   Output: outputs/storyboard/{thread_id}/animated/")
+
+    # Run image generation
+    result = generate_animated_character_images(state)
+
+    if result.get("error"):
+        print(f"❌ Error: {result['error']}")
+        sys.exit(1)
+
+    print(f"\n✅ Generated images in: {result['output_dir']}")
+
+
+if __name__ == "__main__":
+    main()