Funny Pet Image Edits

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "data-designer",
# ]
# ///
"""Funny Pet Image Editing Recipe
Generate a base synthetic dog or cat image, then use image-to-image generation
to make the same pet scene funnier while preserving the pet's identity.
Use this as a playful example of text-to-image followed by image-to-image:
generate a controlled reference image, edit it with additional sampled
conditions, then keep both the image and metadata for visual QA, judge
development, creative review, demos, and model-capability exploration.
Prerequisites:
    - An image-generation provider key for a model that supports image-to-image
      editing through the chat-completions route. The defaults use OpenRouter
      and Gemini 3.1 Flash Image Preview, so set OPENROUTER_API_KEY before running.
Run:
    uv run funny_pet_image_edits.py --num-records 10
"""
from __future__ import annotations
import argparse
from pathlib import Path
import data_designer.config as dd
from data_designer.interface import DataDesigner, DatasetCreationResults
DEFAULT_MODEL_PROVIDER = "openrouter"
DEFAULT_MODEL_ID = "google/gemini-3.1-flash-image-preview"
DEFAULT_MODEL_ALIAS = "funny-pet-image-model"
def build_model_configs(
    *,
    model_provider: str,
    model_id: str,
    model_alias: str,
    image_size: str,
    aspect_ratio: str,
    max_parallel_requests: int,
) -> list[dd.ModelConfig]:
    """Build an image model config for text-to-image and image-to-image generation."""
    return [
        dd.ModelConfig(
            alias=model_alias,
            model=model_id,
            provider=model_provider,
            inference_parameters=dd.ImageInferenceParams(
                extra_body={
                    "modalities": ["image", "text"],
                    "image_config": {
                        "aspect_ratio": aspect_ratio,
                        "image_size": image_size,
                    },
                },
                max_parallel_requests=max_parallel_requests,
            ),
            skip_health_check=True,
        )
    ]
def add_category(config_builder: dd.DataDesignerConfigBuilder, name: str, values: list[str]) -> None:
    """Add a categorical sampler column."""
    config_builder.add_column(
        dd.SamplerColumnConfig(
            name=name,
            sampler_type=dd.SamplerType.CATEGORY,
            params=dd.CategorySamplerParams(values=values),
        )
    )
def add_visual_variation_id(config_builder: dd.DataDesignerConfigBuilder) -> None:
    """Add a unique row-level key that discourages duplicate image generations."""
    config_builder.add_column(
        dd.SamplerColumnConfig(
            name="visual_variation_id",
            sampler_type=dd.SamplerType.UUID,
            params=dd.UUIDSamplerParams(prefix="pet-", short_form=True),
        )
    )
def build_config(
    *,
    model_provider: str = DEFAULT_MODEL_PROVIDER,
    model_id: str = DEFAULT_MODEL_ID,
    model_alias: str = DEFAULT_MODEL_ALIAS,
    image_size: str = "1K",
    aspect_ratio: str = "4:3",
    max_parallel_requests: int = 10,
) -> dd.DataDesignerConfigBuilder:
    """Build a funny pet text-to-image plus image-to-image pipeline."""
    model_configs = build_model_configs(
        model_provider=model_provider,
        model_id=model_id,
        model_alias=model_alias,
        image_size=image_size,
        aspect_ratio=aspect_ratio,
        max_parallel_requests=max_parallel_requests,
    )
    config_builder = dd.DataDesignerConfigBuilder(model_configs=model_configs)
    add_visual_variation_id(config_builder)
    add_category(
        config_builder,
        "pet_type",
        [
            "dog",
            "cat",
        ],
    )
    config_builder.add_column(
        dd.SamplerColumnConfig(
            name="pet_breed",
            sampler_type=dd.SamplerType.SUBCATEGORY,
            params=dd.SubcategorySamplerParams(
                category="pet_type",
                values={
                    "dog": [
                        "German shepherd",
                        "golden retriever",
                        "Pembroke Welsh corgi",
                        "French bulldog",
                        "Shih Tzu",
                        "beagle",
                        "border collie",
                        "mixed-breed terrier",
                    ],
                    "cat": [
                        "orange tabby",
                        "black-and-white tuxedo cat",
                        "gray tabby",
                        "calico",
                        "Siamese cat",
                        "Maine Coon",
                        "British shorthair",
                        "domestic longhair",
                    ],
                },
            ),
        )
    )
    config_builder.add_column(
        dd.SamplerColumnConfig(
            name="pet_age",
            sampler_type=dd.SamplerType.SUBCATEGORY,
            params=dd.SubcategorySamplerParams(
                category="pet_type",
                values={
                    "dog": [
                        "puppy, 6 to 18 months old",
                        "young adult dog, 1 to 4 years old",
                        "adult dog, 4 to 7 years old",
                        "senior dog, 8 years or older",
                    ],
                    "cat": [
                        "kitten, 4 to 12 months old",
                        "young adult cat, 1 to 4 years old",
                        "adult cat, 4 to 10 years old",
                        "senior cat, 11 years or older",
                    ],
                },
            ),
        )
    )
    add_category(
        config_builder,
        "base_activity",
        [
            "sitting proudly at a small table",
            "peeking over the edge of a sofa",
            "standing on a kitchen chair",
            "posing beside a cardboard box",
            "lounging on a soft rug",
            "looking directly at the camera with dramatic seriousness",
            "balanced calmly beside a pile of toys",
        ],
    )
    add_category(
        config_builder,
        "base_setting",
        [
            "sunny living room",
            "cozy home office",
            "tidy kitchen corner",
            "soft studio backdrop",
            "laundry room with folded towels",
            "small apartment balcony with plants",
            "quiet reading nook",
        ],
    )
    add_category(
        config_builder,
        "pet_expression",
        [
            "deeply serious expression",
            "wide-eyed confused expression",
            "proud little smirk",
            "sleepy but determined expression",
            "mildly offended expression",
            "curious head tilt",
        ],
    )
    add_category(
        config_builder,
        "base_photo_style",
        [
            "natural phone photo with soft daylight",
            "clean studio portrait with gentle shadows",
            "warm editorial pet portrait",
            "slightly low-angle comedic portrait",
            "documentary-style candid photo",
        ],
    )
    add_category(
        config_builder,
        "comedy_edit_goal",
        [
            "stage the pet as a tiny orchestra conductor for squeaky toys",
            "stage the pet as a very serious chef inspecting a tiny bowl",
            "stage the pet as a cardboard-spaceship pilot with abstract controls",
            "stage the pet as a detective following a harmless trail of snack crumbs",
            "stage the pet as a living-room sports champion with a tiny trophy",
            "stage the pet as a tiny gardener supervising toy plants",
            "stage the pet as a blanket-cape superhero in a cozy room",
            "stage the pet as a toy stage performer under a tiny spotlight",
        ],
    )
    add_category(
        config_builder,
        "funny_prop",
        [
            "tiny oversized glasses",
            "miniature necktie",
            "small chef hat",
            "toy conductor baton",
            "miniature trophy with no writing",
            "blank tiny clipboard with no writing",
            "cardboard rocket dashboard with colored circles only",
            "toy magnifying glass",
            "small paper crown",
            "tiny blanket cape",
        ],
    )
    add_category(
        config_builder,
        "scene_escalation",
        [
            "add a neatly arranged set of miniature props around the pet",
            "add a playful spotlight and dramatic shadows",
            "add a tiny stage setup made from household objects",
            "add confetti-like paper shapes on the floor",
            "add a pretend control panel made only of colored circles and blank buttons",
            "add an audience of plush toys in the background",
            "add a whimsical but tidy tabletop set",
            "add toy vegetables, an empty bowl, and a tiny spoon",
            "add squeaky toys arranged like an orchestra",
        ],
    )
    add_category(
        config_builder,
        "humor_style",
        [
            "deadpan absurdity",
            "cozy wholesome comedy",
            "overly dramatic tiny-professional energy",
            "gentle visual slapstick without distress",
            "storybook-level silliness",
        ],
    )
    config_builder.add_column(
        dd.ImageColumnConfig(
            name="base_pet_image",
            prompt=BASE_PET_IMAGE_PROMPT,
            model_alias=model_alias,
        )
    )
    config_builder.add_column(
        dd.ImageColumnConfig(
            name="funny_pet_image",
            prompt=FUNNY_PET_EDIT_PROMPT,
            model_alias=model_alias,
            multi_modal_context=[dd.ImageContext(column_name="base_pet_image")],
        )
    )
    return config_builder
def create_dataset(
    config_builder: dd.DataDesignerConfigBuilder,
    *,
    num_records: int,
    dataset_name: str,
    artifact_path: Path | str | None = None,
) -> DatasetCreationResults:
    data_designer = DataDesigner(artifact_path=artifact_path)
    data_designer.validate(config_builder)
    return data_designer.create(config_builder, num_records=num_records, dataset_name=dataset_name)
BASE_PET_IMAGE_PROMPT = """\
Create a realistic synthetic pet photo.
Image requirements:
- Visual variation ID, for internal diversity only: {{ visual_variation_id }}
- Pet type: {{ pet_type }}
- Pet breed: {{ pet_breed }}
- Pet age: {{ pet_age }}
- Base activity: {{ base_activity }}
- Base setting: {{ base_setting }}
- Pet expression: {{ pet_expression }}
- Photo style: {{ base_photo_style }}
Show exactly one healthy, comfortable {{ pet_age }} {{ pet_breed }}. The pet
should be the clear subject, fully visible enough to edit later, and safely
posed in a harmless indoor or domestic setting. Use realistic fur, eyes,
proportions, lighting, shadows, and background details appropriate for the pet
type, breed, and age. Do not include text overlays, real brand logos,
watermarks, captions, speech bubbles, unsafe handling, costumes that restrict
movement, or distressed expressions. Generate exactly one final image for this
row. Use the visual variation ID only as an internal diversity key; never
render it as text.
"""
FUNNY_PET_EDIT_PROMPT = """\
Edit the provided pet image to make the scene funnier while preserving the same pet.
Edit requirements:
- Visual variation ID, for internal diversity only: {{ visual_variation_id }}
- Comedy edit goal: {{ comedy_edit_goal }}
- Funny prop: {{ funny_prop }}
- Scene escalation: {{ scene_escalation }}
- Humor style: {{ humor_style }}
Preserve the same pet identity from the reference image: same species, fur
color, markings, face, body size, expression family, and core pose. Keep the
pet safe, comfortable, healthy, and not distressed. Add playful props,
background details, or scene context that make the image funnier, but keep the
result as one coherent photo-like image rather than a collage.
Do not change the pet into a different animal, add extra pets, add humans,
add speech bubbles, add readable text, add letters, add numbers, add real brand
logos, add watermarks, show unsafe handling, show distress, or make the prop
appear tight, restrictive, or uncomfortable. If papers, signs, screens, labels,
chalkboards, books, control panels, or trophies appear, they must be blank or
use abstract colored shapes only. Generate exactly one final edited image for
this row. Use the visual variation ID only as an internal diversity key; never
render it as text.
"""
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate funny pet image edits.")
    parser.add_argument("--num-records", type=int, default=10, help="Number of funny pet image rows to generate.")
    parser.add_argument("--dataset-name", default="funny-pet-image-edits", help="Output dataset name.")
    parser.add_argument("--artifact-path", type=Path, default=None, help="Optional Data Designer artifact directory.")
    parser.add_argument("--model-provider", default=DEFAULT_MODEL_PROVIDER, help="Image model provider name.")
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID, help="Provider model ID.")
    parser.add_argument("--model-alias", default=DEFAULT_MODEL_ALIAS, help="Alias used by image columns.")
    parser.add_argument("--image-size", default="1K", help="OpenRouter image size tier, such as 1K, 2K, or 4K.")
    parser.add_argument("--aspect-ratio", default="4:3", help="Provider-specific aspect ratio value.")
    parser.add_argument("--max-parallel-requests", type=int, default=10, help="Maximum parallel image requests.")
    return parser.parse_args()
def main() -> None:
    args = parse_args()
    config_builder = build_config(
        model_provider=args.model_provider,
        model_id=args.model_id,
        model_alias=args.model_alias,
        image_size=args.image_size,
        aspect_ratio=args.aspect_ratio,
        max_parallel_requests=args.max_parallel_requests,
    )
    results = create_dataset(
        config_builder,
        num_records=args.num_records,
        dataset_name=args.dataset_name,
        artifact_path=args.artifact_path,
    )
    dataset = results.load_dataset()
    print(f"Generated {len(dataset)} funny pet image-edit rows.")
    print(f"Dataset artifacts: {results.artifact_storage.base_dataset_path}")
if __name__ == "__main__":
    main()