1
import base64
2
import os
3
from pathlib import Path
4
 
5
from IPython.display import Image as IPImage
6
from IPython.display import display
7
 
8
import data_designer.config as dd
9
from data_designer.interface import DataDesigner
10
 

1
data_designer = DataDesigner()
2
 

1
MODEL_PROVIDER = "openrouter"
2
MODEL_ID = "black-forest-labs/flux.2-pro"
3
MODEL_ALIAS = "image-model"
4
 
5
model_configs = [
6
    dd.ModelConfig(
7
        alias=MODEL_ALIAS,
8
        model=MODEL_ID,
9
        provider=MODEL_PROVIDER,
10
        inference_parameters=dd.ImageInferenceParams(
11
            extra_body={"height": 512, "width": 512},
12
        ),
13
    )
14
]
15
 

1
config_builder = dd.DataDesignerConfigBuilder(model_configs=model_configs)
2
 
3
# 1. Sampler columns for diversity
4
config_builder.add_column(
5
    dd.SamplerColumnConfig(
6
        name="animal",
7
        sampler_type=dd.SamplerType.CATEGORY,
8
        params=dd.CategorySamplerParams(
9
            values=["cat", "dog", "fox", "owl", "rabbit", "panda"],
10
        ),
11
    )
12
)
13
 
14
config_builder.add_column(
15
    dd.SamplerColumnConfig(
16
        name="accessory",
17
        sampler_type=dd.SamplerType.CATEGORY,
18
        params=dd.CategorySamplerParams(
19
            values=[
20
                "a tiny top hat",
21
                "oversized sunglasses",
22
                "a red bow tie",
23
                "a knitted beanie",
24
                "a flower crown",
25
                "a monocle and mustache",
26
                "a pirate hat and eye patch",
27
                "a chef hat",
28
            ],
29
        ),
30
    )
31
)
32
 
33
config_builder.add_column(
34
    dd.SamplerColumnConfig(
35
        name="setting",
36
        sampler_type=dd.SamplerType.CATEGORY,
37
        params=dd.CategorySamplerParams(
38
            values=[
39
                "a cozy living room",
40
                "a sunny park",
41
                "a photo studio with soft lighting",
42
                "a red carpet event",
43
                "a holiday card backdrop with snowflakes",
44
                "a tropical beach at sunset",
45
            ],
46
        ),
47
    )
48
)
49
 
50
config_builder.add_column(
51
    dd.SamplerColumnConfig(
52
        name="art_style",
53
        sampler_type=dd.SamplerType.CATEGORY,
54
        params=dd.CategorySamplerParams(
55
            values=[
56
                "a photorealistic style",
57
                "a Disney Pixar 3D render",
58
                "a watercolor painting",
59
                "a pop art poster",
60
            ],
61
        ),
62
    )
63
)
64
 
65
# 2. Generate animal portrait from text
66
config_builder.add_column(
67
    dd.ImageColumnConfig(
68
        name="animal_portrait",
69
        prompt="A close-up portrait photograph of a {{ animal }} looking at the camera, studio lighting, high quality.",
70
        model_alias=MODEL_ALIAS,
71
    )
72
)
73
 
74
# 3. Edit the generated portrait
75
config_builder.add_column(
76
    dd.ImageColumnConfig(
77
        name="edited_portrait",
78
        prompt=(
79
            "Edit this {{ animal }} portrait photo. "
80
            "Add {{ accessory }} on the animal. "
81
            "Place the {{ animal }} in {{ setting }}. "
82
            "Render the result in {{ art_style }}. "
83
            "Keep the animal's face, expression, and features faithful to the original photo."
84
        ),
85
        model_alias=MODEL_ALIAS,
86
        multi_modal_context=[dd.ImageContext(column_name="animal_portrait")],
87
    )
88
)
89
 
90
data_designer.validate(config_builder)
91
 

1
preview = data_designer.preview(config_builder, num_records=2)
2
 

1
for i in range(len(preview.dataset)):
2
    preview.display_sample_record()
3
 

1
preview.dataset
2
 

1
def display_image(image_value, base_path: Path | None = None) -> None:
2
    """Display an image from base64 (preview mode) or file path (create mode)."""
3
    values = [image_value] if isinstance(image_value, str) else list(image_value)
4
    for value in values:
5
        if base_path is not None:
6
            display(IPImage(filename=str(base_path / value)))
7
        else:
8
            display(IPImage(data=base64.b64decode(value)))
9
 
10
 
11
def display_before_after(row, index: int, base_path: Path | None = None) -> None:
12
    """Display original portrait vs edited version for a single record."""
13
    print(f"\n{'=' * 60}")
14
    print(f"Record {index}: {row['animal']} wearing {row['accessory']}")
15
    print(f"Setting: {row['setting']}, Style: {row['art_style']}")
16
    print(f"{'=' * 60}")
17
 
18
    print("\n📷 Generated portrait:")
19
    display_image(row["animal_portrait"], base_path)
20
 
21
    print("\n🎨 Edited version:")
22
    display_image(row["edited_portrait"], base_path)
23
 

1
for index, row in preview.dataset.iterrows():
2
    display_before_after(row, index)
3
 

1
create_num_records = int(os.environ.get("DATA_DESIGNER_FLUX_2_PRO_CREATE_NUM_RECORDS", "5"))
2
results = data_designer.create(
3
    config_builder,
4
    num_records=create_num_records,
5
    dataset_name="tutorial-6-edited-images",
6
)
7
 

1
dataset = results.load_dataset()
2
dataset.head()
3
 

1
for index, row in dataset.head(10).iterrows():
2
    display_before_after(row, index, base_path=results.artifact_storage.base_dataset_path)
3