Source code for pokemon_blip_kandinsky_prior

from mmengine.dataset import DefaultSampler

from diffengine.datasets import HFDataset
from diffengine.datasets.transforms import CLIPImageProcessor, PackInputs
from diffengine.engine.hooks import PriorSaveHook, VisualizationHook

[docs]train_pipeline = [
    dict(type=CLIPImageProcessor, output_key="img",
         pretrained="kandinsky-community/kandinsky-2-2-prior"),
    dict(type=PackInputs),
]
[docs]train_dataloader = dict(
    batch_size=4,
    num_workers=4,
    dataset=dict(
        type=HFDataset,
        dataset="lambdalabs/pokemon-blip-captions",
        pipeline=train_pipeline),
    sampler=dict(type=DefaultSampler, shuffle=True),
)

[docs]val_dataloader = None
[docs]val_evaluator = None
[docs]test_dataloader = val_dataloader
[docs]test_evaluator = val_evaluator

[docs]custom_hooks = [
    dict(type=VisualizationHook, prompt=["yoda pokemon"] * 4,
         height=512, width=512),
    dict(type=PriorSaveHook),
]