Funny Pet Image Edits

1
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
# SPDX-License-Identifier: Apache-2.0
3
# /// script
4
# requires-python = ">=3.10"
5
# dependencies = [
6
#     "data-designer",
7
# ]
8
# ///
9
"""Funny Pet Image Editing Recipe
10
11
Generate a base synthetic dog or cat image, then use image-to-image generation
12
to make the same pet scene funnier while preserving the pet's identity.
13
14
Use this as a playful example of text-to-image followed by image-to-image:
15
generate a controlled reference image, edit it with additional sampled
16
conditions, then keep both the image and metadata for visual QA, judge
17
development, creative review, demos, and model-capability exploration.
18
19
Prerequisites:
20
    - An image-generation provider key for a model that supports image-to-image
21
      editing through the chat-completions route. The defaults use OpenRouter
22
      and Gemini 3.1 Flash Image Preview, so set OPENROUTER_API_KEY before running.
23
24
Run:
25
    uv run funny_pet_image_edits.py --num-records 10
26
"""
27
28
from __future__ import annotations
29
30
import argparse
31
from pathlib import Path
32
33
import data_designer.config as dd
34
from data_designer.interface import DataDesigner, DatasetCreationResults
35
36
DEFAULT_MODEL_PROVIDER = "openrouter"
37
DEFAULT_MODEL_ID = "google/gemini-3.1-flash-image-preview"
38
DEFAULT_MODEL_ALIAS = "funny-pet-image-model"
39
40
41
def build_model_configs(
42
    *,
43
    model_provider: str,
44
    model_id: str,
45
    model_alias: str,
46
    image_size: str,
47
    aspect_ratio: str,
48
    max_parallel_requests: int,
49
) -> list[dd.ModelConfig]:
50
    """Build an image model config for text-to-image and image-to-image generation."""
51
    return [
52
        dd.ModelConfig(
53
            alias=model_alias,
54
            model=model_id,
55
            provider=model_provider,
56
            inference_parameters=dd.ImageInferenceParams(
57
                extra_body={
58
                    "modalities": ["image", "text"],
59
                    "image_config": {
60
                        "aspect_ratio": aspect_ratio,
61
                        "image_size": image_size,
62
                    },
63
                },
64
                max_parallel_requests=max_parallel_requests,
65
            ),
66
            skip_health_check=True,
67
        )
68
    ]
69
70
71
def add_category(config_builder: dd.DataDesignerConfigBuilder, name: str, values: list[str]) -> None:
72
    """Add a categorical sampler column."""
73
    config_builder.add_column(
74
        dd.SamplerColumnConfig(
75
            name=name,
76
            sampler_type=dd.SamplerType.CATEGORY,
77
            params=dd.CategorySamplerParams(values=values),
78
        )
79
    )
80
81
82
def add_visual_variation_id(config_builder: dd.DataDesignerConfigBuilder) -> None:
83
    """Add a unique row-level key that discourages duplicate image generations."""
84
    config_builder.add_column(
85
        dd.SamplerColumnConfig(
86
            name="visual_variation_id",
87
            sampler_type=dd.SamplerType.UUID,
88
            params=dd.UUIDSamplerParams(prefix="pet-", short_form=True),
89
        )
90
    )
91
92
93
def build_config(
94
    *,
95
    model_provider: str = DEFAULT_MODEL_PROVIDER,
96
    model_id: str = DEFAULT_MODEL_ID,
97
    model_alias: str = DEFAULT_MODEL_ALIAS,
98
    image_size: str = "1K",
99
    aspect_ratio: str = "4:3",
100
    max_parallel_requests: int = 10,
101
) -> dd.DataDesignerConfigBuilder:
102
    """Build a funny pet text-to-image plus image-to-image pipeline."""
103
    model_configs = build_model_configs(
104
        model_provider=model_provider,
105
        model_id=model_id,
106
        model_alias=model_alias,
107
        image_size=image_size,
108
        aspect_ratio=aspect_ratio,
109
        max_parallel_requests=max_parallel_requests,
110
    )
111
    config_builder = dd.DataDesignerConfigBuilder(model_configs=model_configs)
112
    add_visual_variation_id(config_builder)
113
114
    add_category(
115
        config_builder,
116
        "pet_type",
117
        [
118
            "dog",
119
            "cat",
120
        ],
121
    )
122
    config_builder.add_column(
123
        dd.SamplerColumnConfig(
124
            name="pet_breed",
125
            sampler_type=dd.SamplerType.SUBCATEGORY,
126
            params=dd.SubcategorySamplerParams(
127
                category="pet_type",
128
                values={
129
                    "dog": [
130
                        "German shepherd",
131
                        "golden retriever",
132
                        "Pembroke Welsh corgi",
133
                        "French bulldog",
134
                        "Shih Tzu",
135
                        "beagle",
136
                        "border collie",
137
                        "mixed-breed terrier",
138
                    ],
139
                    "cat": [
140
                        "orange tabby",
141
                        "black-and-white tuxedo cat",
142
                        "gray tabby",
143
                        "calico",
144
                        "Siamese cat",
145
                        "Maine Coon",
146
                        "British shorthair",
147
                        "domestic longhair",
148
                    ],
149
                },
150
            ),
151
        )
152
    )
153
    config_builder.add_column(
154
        dd.SamplerColumnConfig(
155
            name="pet_age",
156
            sampler_type=dd.SamplerType.SUBCATEGORY,
157
            params=dd.SubcategorySamplerParams(
158
                category="pet_type",
159
                values={
160
                    "dog": [
161
                        "puppy, 6 to 18 months old",
162
                        "young adult dog, 1 to 4 years old",
163
                        "adult dog, 4 to 7 years old",
164
                        "senior dog, 8 years or older",
165
                    ],
166
                    "cat": [
167
                        "kitten, 4 to 12 months old",
168
                        "young adult cat, 1 to 4 years old",
169
                        "adult cat, 4 to 10 years old",
170
                        "senior cat, 11 years or older",
171
                    ],
172
                },
173
            ),
174
        )
175
    )
176
    add_category(
177
        config_builder,
178
        "base_activity",
179
        [
180
            "sitting proudly at a small table",
181
            "peeking over the edge of a sofa",
182
            "standing on a kitchen chair",
183
            "posing beside a cardboard box",
184
            "lounging on a soft rug",
185
            "looking directly at the camera with dramatic seriousness",
186
            "balanced calmly beside a pile of toys",
187
        ],
188
    )
189
    add_category(
190
        config_builder,
191
        "base_setting",
192
        [
193
            "sunny living room",
194
            "cozy home office",
195
            "tidy kitchen corner",
196
            "soft studio backdrop",
197
            "laundry room with folded towels",
198
            "small apartment balcony with plants",
199
            "quiet reading nook",
200
        ],
201
    )
202
    add_category(
203
        config_builder,
204
        "pet_expression",
205
        [
206
            "deeply serious expression",
207
            "wide-eyed confused expression",
208
            "proud little smirk",
209
            "sleepy but determined expression",
210
            "mildly offended expression",
211
            "curious head tilt",
212
        ],
213
    )
214
    add_category(
215
        config_builder,
216
        "base_photo_style",
217
        [
218
            "natural phone photo with soft daylight",
219
            "clean studio portrait with gentle shadows",
220
            "warm editorial pet portrait",
221
            "slightly low-angle comedic portrait",
222
            "documentary-style candid photo",
223
        ],
224
    )
225
    add_category(
226
        config_builder,
227
        "comedy_edit_goal",
228
        [
229
            "stage the pet as a tiny orchestra conductor for squeaky toys",
230
            "stage the pet as a very serious chef inspecting a tiny bowl",
231
            "stage the pet as a cardboard-spaceship pilot with abstract controls",
232
            "stage the pet as a detective following a harmless trail of snack crumbs",
233
            "stage the pet as a living-room sports champion with a tiny trophy",
234
            "stage the pet as a tiny gardener supervising toy plants",
235
            "stage the pet as a blanket-cape superhero in a cozy room",
236
            "stage the pet as a toy stage performer under a tiny spotlight",
237
        ],
238
    )
239
    add_category(
240
        config_builder,
241
        "funny_prop",
242
        [
243
            "tiny oversized glasses",
244
            "miniature necktie",
245
            "small chef hat",
246
            "toy conductor baton",
247
            "miniature trophy with no writing",
248
            "blank tiny clipboard with no writing",
249
            "cardboard rocket dashboard with colored circles only",
250
            "toy magnifying glass",
251
            "small paper crown",
252
            "tiny blanket cape",
253
        ],
254
    )
255
    add_category(
256
        config_builder,
257
        "scene_escalation",
258
        [
259
            "add a neatly arranged set of miniature props around the pet",
260
            "add a playful spotlight and dramatic shadows",
261
            "add a tiny stage setup made from household objects",
262
            "add confetti-like paper shapes on the floor",
263
            "add a pretend control panel made only of colored circles and blank buttons",
264
            "add an audience of plush toys in the background",
265
            "add a whimsical but tidy tabletop set",
266
            "add toy vegetables, an empty bowl, and a tiny spoon",
267
            "add squeaky toys arranged like an orchestra",
268
        ],
269
    )
270
    add_category(
271
        config_builder,
272
        "humor_style",
273
        [
274
            "deadpan absurdity",
275
            "cozy wholesome comedy",
276
            "overly dramatic tiny-professional energy",
277
            "gentle visual slapstick without distress",
278
            "storybook-level silliness",
279
        ],
280
    )
281
282
    config_builder.add_column(
283
        dd.ImageColumnConfig(
284
            name="base_pet_image",
285
            prompt=BASE_PET_IMAGE_PROMPT,
286
            model_alias=model_alias,
287
        )
288
    )
289
290
    config_builder.add_column(
291
        dd.ImageColumnConfig(
292
            name="funny_pet_image",
293
            prompt=FUNNY_PET_EDIT_PROMPT,
294
            model_alias=model_alias,
295
            multi_modal_context=[dd.ImageContext(column_name="base_pet_image")],
296
        )
297
    )
298
299
    return config_builder
300
301
302
def create_dataset(
303
    config_builder: dd.DataDesignerConfigBuilder,
304
    *,
305
    num_records: int,
306
    dataset_name: str,
307
    artifact_path: Path | str | None = None,
308
) -> DatasetCreationResults:
309
    data_designer = DataDesigner(artifact_path=artifact_path)
310
    data_designer.validate(config_builder)
311
    return data_designer.create(config_builder, num_records=num_records, dataset_name=dataset_name)
312
313
314
BASE_PET_IMAGE_PROMPT = """\
315
Create a realistic synthetic pet photo.
316
317
Image requirements:
318
- Visual variation ID, for internal diversity only: {{ visual_variation_id }}
319
- Pet type: {{ pet_type }}
320
- Pet breed: {{ pet_breed }}
321
- Pet age: {{ pet_age }}
322
- Base activity: {{ base_activity }}
323
- Base setting: {{ base_setting }}
324
- Pet expression: {{ pet_expression }}
325
- Photo style: {{ base_photo_style }}
326
327
Show exactly one healthy, comfortable {{ pet_age }} {{ pet_breed }}. The pet
328
should be the clear subject, fully visible enough to edit later, and safely
329
posed in a harmless indoor or domestic setting. Use realistic fur, eyes,
330
proportions, lighting, shadows, and background details appropriate for the pet
331
type, breed, and age. Do not include text overlays, real brand logos,
332
watermarks, captions, speech bubbles, unsafe handling, costumes that restrict
333
movement, or distressed expressions. Generate exactly one final image for this
334
row. Use the visual variation ID only as an internal diversity key; never
335
render it as text.
336
"""
337
338
339
FUNNY_PET_EDIT_PROMPT = """\
340
Edit the provided pet image to make the scene funnier while preserving the same pet.
341
342
Edit requirements:
343
- Visual variation ID, for internal diversity only: {{ visual_variation_id }}
344
- Comedy edit goal: {{ comedy_edit_goal }}
345
- Funny prop: {{ funny_prop }}
346
- Scene escalation: {{ scene_escalation }}
347
- Humor style: {{ humor_style }}
348
349
Preserve the same pet identity from the reference image: same species, fur
350
color, markings, face, body size, expression family, and core pose. Keep the
351
pet safe, comfortable, healthy, and not distressed. Add playful props,
352
background details, or scene context that make the image funnier, but keep the
353
result as one coherent photo-like image rather than a collage.
354
355
Do not change the pet into a different animal, add extra pets, add humans,
356
add speech bubbles, add readable text, add letters, add numbers, add real brand
357
logos, add watermarks, show unsafe handling, show distress, or make the prop
358
appear tight, restrictive, or uncomfortable. If papers, signs, screens, labels,
359
chalkboards, books, control panels, or trophies appear, they must be blank or
360
use abstract colored shapes only. Generate exactly one final edited image for
361
this row. Use the visual variation ID only as an internal diversity key; never
362
render it as text.
363
"""
364
365
366
def parse_args() -> argparse.Namespace:
367
    parser = argparse.ArgumentParser(description="Generate funny pet image edits.")
368
    parser.add_argument("--num-records", type=int, default=10, help="Number of funny pet image rows to generate.")
369
    parser.add_argument("--dataset-name", default="funny-pet-image-edits", help="Output dataset name.")
370
    parser.add_argument("--artifact-path", type=Path, default=None, help="Optional Data Designer artifact directory.")
371
    parser.add_argument("--model-provider", default=DEFAULT_MODEL_PROVIDER, help="Image model provider name.")
372
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID, help="Provider model ID.")
373
    parser.add_argument("--model-alias", default=DEFAULT_MODEL_ALIAS, help="Alias used by image columns.")
374
    parser.add_argument("--image-size", default="1K", help="OpenRouter image size tier, such as 1K, 2K, or 4K.")
375
    parser.add_argument("--aspect-ratio", default="4:3", help="Provider-specific aspect ratio value.")
376
    parser.add_argument("--max-parallel-requests", type=int, default=10, help="Maximum parallel image requests.")
377
    return parser.parse_args()
378
379
380
def main() -> None:
381
    args = parse_args()
382
    config_builder = build_config(
383
        model_provider=args.model_provider,
384
        model_id=args.model_id,
385
        model_alias=args.model_alias,
386
        image_size=args.image_size,
387
        aspect_ratio=args.aspect_ratio,
388
        max_parallel_requests=args.max_parallel_requests,
389
    )
390
    results = create_dataset(
391
        config_builder,
392
        num_records=args.num_records,
393
        dataset_name=args.dataset_name,
394
        artifact_path=args.artifact_path,
395
    )
396
    dataset = results.load_dataset()
397
    print(f"Generated {len(dataset)} funny pet image-edit rows.")
398
    print(f"Dataset artifacts: {results.artifact_storage.base_dataset_path}")
399
400
401
if __name__ == "__main__":
402
    main()