Funny Pet Image Edits
Download Recipe
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()