Upload folder using huggingface_hub (part 4)

examples/ideogram4/model_training/lora/Ideogram-4-fp8.sh

@@ -0,0 +1,18 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ideogram4/Ideogram-4-fp8/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/ideogram4/model_training/train.py \
+    --learning_rate 1e-4 \
+    --num_epochs 5 \
+    --lora_rank 32 \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --dataset_base_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-fp8" \
+    --dataset_metadata_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-fp8/metadata.json" \
+    --model_id_with_origin_paths "ideogram-ai/ideogram-4-fp8:transformer/diffusion_pytorch_model.safetensors,ideogram-ai/ideogram-4-fp8:text_encoder/model.safetensors,ideogram-ai/ideogram-4-fp8:vae/diffusion_pytorch_model.safetensors" \
+    --lora_base_model "dit" \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --max_pixels 1048576 \
+    --dataset_repeat 100 \
+    --output_path "./models/train/Ideogram-4-fp8_lora" \
+    --lora_target_modules "attention.qkv,attention.o,feed_forward.w1,feed_forward.w2,feed_forward.w3,adaln_modulation" \
+    --data_file_keys "image"

examples/ideogram4/model_training/lora/Ideogram-4-nf4.sh

@@ -0,0 +1,18 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ideogram4/Ideogram-4-nf4/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/ideogram4/model_training/train.py \
+    --learning_rate 1e-4 \
+    --num_epochs 5 \
+    --lora_rank 32 \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --dataset_base_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-nf4" \
+    --dataset_metadata_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-nf4/metadata.json" \
+    --model_id_with_origin_paths "ideogram-ai/ideogram-4-nf4:transformer/diffusion_pytorch_model.safetensors,ideogram-ai/ideogram-4-nf4:text_encoder/model.safetensors,ideogram-ai/ideogram-4-nf4:vae/diffusion_pytorch_model.safetensors" \
+    --lora_base_model "dit" \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --max_pixels 1048576 \
+    --dataset_repeat 100 \
+    --output_path "./models/train/Ideogram-4-nf4_lora" \
+    --lora_target_modules "attention.qkv,attention.o,feed_forward.w1,feed_forward.w2,feed_forward.w3,adaln_modulation" \
+    --data_file_keys "image"

examples/ideogram4/model_training/special/split_training/Ideogram-4-bf16-repackage.sh

@@ -0,0 +1,40 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ideogram4/Ideogram-4-bf16-repackage/*" --local_dir ./data/diffsynth_example_dataset
+
+# Stage 1: cache deterministic preprocessing outputs.
+accelerate launch examples/ideogram4/model_training/train.py \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --dataset_base_path ./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage \
+  --dataset_metadata_path ./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage/metadata.json \
+  --model_id_with_origin_paths DiffSynth-Studio/ideogram-4-bf16-repackage:transformer/diffusion_pytorch_model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:text_encoder/model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:vae/diffusion_pytorch_model.safetensors \
+  --lora_base_model dit \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --max_pixels 1048576 \
+  --dataset_repeat 1 \
+  --output_path ./models/train/Ideogram-4-bf16-repackage_split_cache \
+  --lora_target_modules attention.qkv,attention.o,feed_forward.w1,feed_forward.w2,feed_forward.w3,adaln_modulation \
+  --data_file_keys image \
+  --offload_models DiffSynth-Studio/ideogram-4-bf16-repackage:transformer/diffusion_pytorch_model.safetensors \
+  --task sft:data_process
+
+# Stage 2: train LoRA from the cached dataset.
+accelerate launch examples/ideogram4/model_training/train.py \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --dataset_base_path ./models/train/Ideogram-4-bf16-repackage_split_cache \
+  --model_id_with_origin_paths DiffSynth-Studio/ideogram-4-bf16-repackage:transformer/diffusion_pytorch_model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:text_encoder/model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:vae/diffusion_pytorch_model.safetensors \
+  --lora_base_model dit \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --max_pixels 1048576 \
+  --dataset_repeat 100 \
+  --output_path ./models/train/Ideogram-4-bf16-repackage_split \
+  --lora_target_modules attention.qkv,attention.o,feed_forward.w1,feed_forward.w2,feed_forward.w3,adaln_modulation \
+  --data_file_keys image \
+  --offload_models DiffSynth-Studio/ideogram-4-bf16-repackage:text_encoder/model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:vae/diffusion_pytorch_model.safetensors \
+  --task sft:train

examples/ideogram4/model_training/special/split_training/validate.py

@@ -0,0 +1,21 @@
+from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline
+from diffsynth.core import ModelConfig
+import torch
+
+pipe = Ideogram4Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"),
+)
+pipe.load_lora(pipe.dit, './models/train/Ideogram-4-bf16-repackage_split/epoch-4.safetensors', alpha=1)
+# pipe.load_lora(pipe.dit_uncond, './models/train/Ideogram-4-bf16-repackage_split/epoch-4.safetensors', alpha=1)
+
+prompt = "{\"high_level_description\":\"A close-up photograph of a happy Pembroke Welsh Corgi sitting on a concrete wall, panting with its tongue out, set against a backdrop of blurred pink cherry blossoms and blue sky.\",\"style_description\":{\"aesthetics\":\"joyful, vibrant, spring-like, cute, energetic\",\"lighting\":\"bright natural daylight, soft diffuse sunlight, shallow depth of field\",\"photo\":\"85mm lens, f/2.0, bokeh background, sharp focus on dog face\",\"medium\":\"photograph\",\"color_palette\":[\"#E2725B\",\"#FFFFFF\",\"#FFB7C5\",\"#87CEEB\",\"#A9A9A9\"]},\"compositional_deconstruction\":{\"background\":\"Softly blurred background of pink cherry blossom branches against a pale blue sky. The bokeh effect creates a dreamy spring atmosphere. The background is out of focus to highlight the sharp details of the dog in the foreground.\",\"elements\":[{\"type\":\"obj\",\"bbox\":[150,200,900,850],\"desc\":\"A Pembroke Welsh Corgi with fluffy orange and white fur. Its mouth is open, panting with a pink tongue hanging out, expression is happy and excited. Ears are perked up. Sharp focus on the face and eyes.\"},{\"type\":\"obj\",\"bbox\":[850,0,1000,1000],\"desc\":\"A grey concrete wall or ledge at the bottom of the frame. The dog's front paws are resting near the edge. Rough texture.\"}]}}"
+image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=0)
+image.save('split_training_Ideogram-4.jpg')

examples/ideogram4/model_training/train.py

@@ -0,0 +1,132 @@
+import torch, os, argparse, accelerate
+from diffsynth.core import UnifiedDataset
+from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline, ModelConfig
+from diffsynth.diffusion import *
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class Ideogram4TrainingModule(DiffusionTrainingModule):
+    def __init__(
+        self,
+        model_paths=None, model_id_with_origin_paths=None,
+        tokenizer_path=None,
+        trainable_models=None,
+        lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None,
+        preset_lora_path=None, preset_lora_model=None,
+        use_gradient_checkpointing=True,
+        use_gradient_checkpointing_offload=False,
+        extra_inputs=None,
+        fp8_models=None,
+        offload_models=None,
+        quant_options=None,
+        device="cpu",
+        task="sft",
+    ):
+        super().__init__()
+        model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, quant_options=quant_options, device=device)
+        tokenizer_config = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"))
+        self.pipe = Ideogram4Pipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config)
+        self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model)
+
+        self.switch_pipe_to_training_mode(
+            self.pipe, trainable_models,
+            lora_base_model, lora_target_modules, lora_rank, lora_checkpoint,
+            preset_lora_path, preset_lora_model,
+            task=task,
+        )
+
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
+        self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else []
+        self.fp8_models = fp8_models
+        self.task = task
+        self.task_to_loss = {
+            "sft:data_process": lambda pipe, *args: args,
+            "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+            "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+        }
+
+    def get_pipeline_inputs(self, data):
+        inputs_posi = {"prompt": data["prompt"]}
+        inputs_nega = {"negative_prompt": ""}
+        inputs_shared = {
+            "input_image": data["image"],
+            "height": data["image"].size[1],
+            "width": data["image"].size[0],
+            "cfg_scale": 1,
+            "rand_device": self.pipe.device,
+            "use_gradient_checkpointing": self.use_gradient_checkpointing,
+            "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload,
+        }
+        inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared)
+        return inputs_shared, inputs_posi, inputs_nega
+
+    def forward(self, data, inputs=None):
+        if inputs is None: inputs = self.get_pipeline_inputs(data)
+        inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype)
+        for unit in self.pipe.units:
+            inputs = self.pipe.unit_runner(unit, self.pipe, *inputs)
+        loss = self.task_to_loss[self.task](self.pipe, *inputs)
+        return loss
+
+
+def ideogram4_parser():
+    parser = argparse.ArgumentParser(description="Ideogram4 training.")
+    parser = add_general_config(parser)
+    parser = add_image_size_config(parser)
+    parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.")
+    return parser
+
+
+if __name__ == "__main__":
+    parser = ideogram4_parser()
+    args = parser.parse_args()
+    accelerator = accelerate.Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)],
+    )
+    dataset = UnifiedDataset(
+        base_path=args.dataset_base_path,
+        metadata_path=args.dataset_metadata_path,
+        repeat=args.dataset_repeat,
+        data_file_keys=args.data_file_keys.split(","),
+        main_data_operator=UnifiedDataset.default_image_operator(
+            base_path=args.dataset_base_path,
+            max_pixels=args.max_pixels,
+            height=args.height,
+            width=args.width,
+            height_division_factor=16,
+            width_division_factor=16,
+        ),
+    )
+    model = Ideogram4TrainingModule(
+        model_paths=args.model_paths,
+        model_id_with_origin_paths=args.model_id_with_origin_paths,
+        tokenizer_path=args.tokenizer_path,
+        trainable_models=args.trainable_models,
+        lora_base_model=args.lora_base_model,
+        lora_target_modules=args.lora_target_modules,
+        lora_rank=args.lora_rank,
+        lora_checkpoint=args.lora_checkpoint,
+        preset_lora_path=args.preset_lora_path,
+        preset_lora_model=args.preset_lora_model,
+        use_gradient_checkpointing=args.use_gradient_checkpointing,
+        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
+        extra_inputs=args.extra_inputs,
+        fp8_models=args.fp8_models,
+        offload_models=args.offload_models,
+        quant_options=args.quant_options,
+        task=args.task,
+        device=accelerator.device,
+    )
+    model_logger = ModelLogger(
+        args.output_path,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        enable_csv_log=args.enable_csv_log,
+    )
+    launcher_map = {
+        "sft:data_process": launch_data_process_task,
+        "sft": launch_training_task,
+        "sft:train": launch_training_task,
+    }
+    launcher_map[args.task](accelerator, dataset, model, model_logger, args=args)

examples/ideogram4/model_training/validate_lora/Ideogram-4-bf16-repackage.py

@@ -0,0 +1,21 @@
+from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline
+from diffsynth.core import ModelConfig
+import torch
+
+pipe = Ideogram4Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"),
+)
+pipe.load_lora(pipe.dit, "models/train/Ideogram-4-bf16-repackage_lora/epoch-4.safetensors", alpha=1)
+pipe.load_lora(pipe.dit_uncond, "models/train/Ideogram-4-bf16-repackage_lora/epoch-4.safetensors", alpha=1)
+
+prompt = "{\"high_level_description\":\"A close-up photograph of a happy Pembroke Welsh Corgi sitting on a concrete wall, panting with its tongue out, set against a backdrop of blurred pink cherry blossoms and blue sky.\",\"style_description\":{\"aesthetics\":\"joyful, vibrant, spring-like, cute, energetic\",\"lighting\":\"bright natural daylight, soft diffuse sunlight, shallow depth of field\",\"photo\":\"85mm lens, f/2.0, bokeh background, sharp focus on dog face\",\"medium\":\"photograph\",\"color_palette\":[\"#E2725B\",\"#FFFFFF\",\"#FFB7C5\",\"#87CEEB\",\"#A9A9A9\"]},\"compositional_deconstruction\":{\"background\":\"Softly blurred background of pink cherry blossom branches against a pale blue sky. The bokeh effect creates a dreamy spring atmosphere. The background is out of focus to highlight the sharp details of the dog in the foreground.\",\"elements\":[{\"type\":\"obj\",\"bbox\":[150,200,900,850],\"desc\":\"A Pembroke Welsh Corgi with fluffy orange and white fur. Its mouth is open, panting with a pink tongue hanging out, expression is happy and excited. Ears are perked up. Sharp focus on the face and eyes.\"},{\"type\":\"obj\",\"bbox\":[850,0,1000,1000],\"desc\":\"A grey concrete wall or ledge at the bottom of the frame. The dog's front paws are resting near the edge. Rough texture.\"}]}}"
+image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=0)
+image.save("image_lora.jpg")

examples/ideogram4/model_training/validate_lora/Ideogram-4-fp8.py

@@ -0,0 +1,32 @@
+from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline
+from diffsynth.core import ModelConfig
+import torch
+
+vram_config = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cpu",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = Ideogram4Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+pipe.load_lora(pipe.dit, "models/train/Ideogram-4-fp8_lora/epoch-4.safetensors", alpha=1)
+pipe.load_lora(pipe.dit_uncond, "models/train/Ideogram-4-fp8_lora/epoch-4.safetensors", alpha=1)
+
+prompt = "{\"high_level_description\":\"A close-up photograph of a happy Pembroke Welsh Corgi sitting on a concrete wall, panting with its tongue out, set against a backdrop of blurred pink cherry blossoms and blue sky.\",\"style_description\":{\"aesthetics\":\"joyful, vibrant, spring-like, cute, energetic\",\"lighting\":\"bright natural daylight, soft diffuse sunlight, shallow depth of field\",\"photo\":\"85mm lens, f/2.0, bokeh background, sharp focus on dog face\",\"medium\":\"photograph\",\"color_palette\":[\"#E2725B\",\"#FFFFFF\",\"#FFB7C5\",\"#87CEEB\",\"#A9A9A9\"]},\"compositional_deconstruction\":{\"background\":\"Softly blurred background of pink cherry blossom branches against a pale blue sky. The bokeh effect creates a dreamy spring atmosphere. The background is out of focus to highlight the sharp details of the dog in the foreground.\",\"elements\":[{\"type\":\"obj\",\"bbox\":[150,200,900,850],\"desc\":\"A Pembroke Welsh Corgi with fluffy orange and white fur. Its mouth is open, panting with a pink tongue hanging out, expression is happy and excited. Ears are perked up. Sharp focus on the face and eyes.\"},{\"type\":\"obj\",\"bbox\":[850,0,1000,1000],\"desc\":\"A grey concrete wall or ledge at the bottom of the frame. The dog's front paws are resting near the edge. Rough texture.\"}]}}"
+image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=0)
+image.save("image_lora.jpg")

examples/ideogram4/model_training/validate_lora/Ideogram-4-nf4.py

@@ -0,0 +1,32 @@
+from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline
+from diffsynth.core import ModelConfig
+import torch
+
+vram_config = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cpu",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = Ideogram4Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="tokenizer/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+pipe.load_lora(pipe.dit, "models/train/Ideogram-4-nf4_lora/epoch-4.safetensors", alpha=1)
+pipe.load_lora(pipe.dit_uncond, "models/train/Ideogram-4-nf4_lora/epoch-4.safetensors", alpha=1)
+
+prompt = "{\"high_level_description\":\"A close-up photograph of a happy Pembroke Welsh Corgi sitting on a concrete wall, panting with its tongue out, set against a backdrop of blurred pink cherry blossoms and blue sky.\",\"style_description\":{\"aesthetics\":\"joyful, vibrant, spring-like, cute, energetic\",\"lighting\":\"bright natural daylight, soft diffuse sunlight, shallow depth of field\",\"photo\":\"85mm lens, f/2.0, bokeh background, sharp focus on dog face\",\"medium\":\"photograph\",\"color_palette\":[\"#E2725B\",\"#FFFFFF\",\"#FFB7C5\",\"#87CEEB\",\"#A9A9A9\"]},\"compositional_deconstruction\":{\"background\":\"Softly blurred background of pink cherry blossom branches against a pale blue sky. The bokeh effect creates a dreamy spring atmosphere. The background is out of focus to highlight the sharp details of the dog in the foreground.\",\"elements\":[{\"type\":\"obj\",\"bbox\":[150,200,900,850],\"desc\":\"A Pembroke Welsh Corgi with fluffy orange and white fur. Its mouth is open, panting with a pink tongue hanging out, expression is happy and excited. Ears are perked up. Sharp focus on the face and eyes.\"},{\"type\":\"obj\",\"bbox\":[850,0,1000,1000],\"desc\":\"A grey concrete wall or ledge at the bottom of the frame. The dog's front paws are resting near the edge. Rough texture.\"}]}}"
+image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=0)
+image.save("image_lora.jpg")

examples/image_quality_metric/aesthetic.py

@@ -0,0 +1,16 @@
+from diffsynth.metrics import AestheticMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+metric = AestheticMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="Aesthetic/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(image)[0]
+print(f"Aesthetic score: {score:.3f}")

examples/image_quality_metric/bioclip.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import BioCLIPMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a photo of Animalia Chordata Mammalia Carnivora Canidae Canis Canis lupus familiaris with common name domestic dog."
+metric = BioCLIPMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="BioCLIPv2/open_clip_model.safetensors"),
+    device="cuda",
+)
+score = metric.compute(prompt, image)[0]
+print(f"BioCLIP score: {score:.3f}")

examples/image_quality_metric/clipscore.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import CLIPMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = CLIPMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="CLIP-ViT-H-14-laion2B-s32B-b79K/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(prompt, image)[0]
+print(f"CLIP score: {score:.3f}")

examples/image_quality_metric/fid.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import FIDMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern=["flux/FLUX.1-dev/*", "flux2/FLUX.2-dev/*"],
+    local_dir="./data/diffsynth_example_dataset",
+)
+metric = FIDMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="FID/model.safetensors"),
+    device="cuda",
+)
+score = metric.compute(
+    "./data/diffsynth_example_dataset/flux/FLUX.1-dev",
+    "./data/diffsynth_example_dataset/flux2/FLUX.2-dev",
+)
+print(f"FID score: {score:.3f}")

examples/image_quality_metric/hpsv2.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import HPSv2Metric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = HPSv2Metric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv2/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(prompt, image)[0]
+print(f"HPSv2 score: {score:.3f}")

examples/image_quality_metric/hpsv3.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import HPSv3Metric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = HPSv3Metric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="HPSv3/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(prompt, image)[0]
+print(f"HPSv3 score: {score:.3f}")

examples/image_quality_metric/image_reward.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import ImageRewardMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = ImageRewardMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="ImageReward/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(prompt, image)[0]
+print(f"ImageReward score: {score:.3f}")

examples/image_quality_metric/lpips.py

@@ -0,0 +1,33 @@
+from diffsynth.metrics import LPIPSMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern=["flux/FLUX.1-dev/*", "flux2/FLUX.2-dev/*"],
+    local_dir="./data/diffsynth_example_dataset",
+)
+
+# net="alex" with LPIPS/alexnet.safetensors (default)
+# For VGG: net="vgg", model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="LPIPS/vgg.safetensors")
+# For SqueezeNet: net="squeeze", model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="LPIPS/squeezenet.safetensors")
+metric = LPIPSMetric.from_pretrained(
+    net="alex",
+    model_config=ModelConfig(
+        model_id="DiffSynth-Studio/ImageMetrics",
+        origin_file_pattern="LPIPS/alexnet.safetensors",
+    ),
+    device="cuda",
+    target_size=512,
+)
+
+score = metric.compute(
+    "./data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg",
+    "./data/diffsynth_example_dataset/flux/FLUX.1-dev/2.jpg",
+)
+print(f"LPIPS score (image vs image): {score:.4f}")
+
+score = metric.compute(
+    "./data/diffsynth_example_dataset/flux/FLUX.1-dev",
+    "./data/diffsynth_example_dataset/flux2/FLUX.2-dev",
+)
+print(f"LPIPS score (dir vs dir): {score:.4f}")

examples/image_quality_metric/pickscore.py

@@ -0,0 +1,17 @@
+from diffsynth.metrics import PickScoreMetric, ModelConfig
+from modelscope import dataset_snapshot_download
+from PIL import Image
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = PickScoreMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="PickScore/model.safetensors"),
+    device="cuda"
+)
+score = metric.compute(prompt, image)[0]
+print(f"PickScore score:: {score:.3f}")

examples/image_quality_metric/qwen_image_bench.py

@@ -0,0 +1,21 @@
+from diffsynth.metrics import ModelConfig, QwenImageBenchMetric
+from modelscope import dataset_snapshot_download
+from PIL import Image
+import json
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = QwenImageBenchMetric.from_pretrained(
+    model_config=ModelConfig(model_id="Qwen/Qwen-Image-Bench", origin_file_pattern="model-*.safetensors"),
+    processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Bench"),
+    device="cuda",
+)
+details = metric.evaluate(prompt, image)[0]
+print(json.dumps(details["level1_scores"], indent=4, ensure_ascii=False))
+print(json.dumps(details["level2_scores"], indent=4, ensure_ascii=False))
+print(json.dumps(details["level3_scores"], indent=4, ensure_ascii=False))

examples/image_quality_metric/unified_reward_2.py

@@ -0,0 +1,19 @@
+from diffsynth.metrics import ModelConfig, UnifiedReward2Metric
+from modelscope import dataset_snapshot_download
+from PIL import Image
+import json
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="flux/FLUX.1-dev/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+image = Image.open("data/diffsynth_example_dataset/flux/FLUX.1-dev/1.jpg").convert("RGB")
+prompt = "a dog"
+metric = UnifiedReward2Metric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-2.0-qwen35-9b/model-*.safetensors"),
+    processor_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-2.0-qwen35-9b/"),
+    device="cuda",
+)
+details = metric.evaluate(prompt, image)[0]
+print(json.dumps(details, indent=4, ensure_ascii=False))

examples/image_quality_metric/unified_reward_edit.py

@@ -0,0 +1,25 @@
+from diffsynth.metrics import ModelConfig, UnifiedRewardEditMetric
+from modelscope import dataset_snapshot_download
+from PIL import Image
+import json
+
+dataset_snapshot_download(
+    "DiffSynth-Studio/diffsynth_example_dataset",
+    allow_file_pattern="qwen_image/Qwen-Image-Edit/*",
+    local_dir="./data/diffsynth_example_dataset",
+)
+source_image = Image.open("data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit/edit/image1.jpg").convert("RGB")
+edited_image_1 = Image.open("data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit/edit/image2.jpg").convert("RGB")
+edited_image_2 = source_image
+instruction = "Change the dress to pink."
+metric = UnifiedRewardEditMetric.from_pretrained(
+    model_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-Edit-qwen3vl-8b/model-*.safetensors"),  
+    processor_config=ModelConfig(model_id="DiffSynth-Studio/ImageMetrics", origin_file_pattern="UnifiedReward-Edit-qwen3vl-8b/"),
+    device="cuda",
+)
+details = metric.evaluate(instruction, [source_image, edited_image_1], task="edit_pointwise_score")[0]
+print(json.dumps(details, indent=4, ensure_ascii=False))
+details = metric.evaluate(instruction, [source_image, edited_image_1, edited_image_2], task="edit_pairwise_rank")[0]
+print(json.dumps(details, indent=4, ensure_ascii=False))
+details = metric.evaluate(instruction, [source_image, edited_image_1, edited_image_2], task="edit_pairwise_score")[0]
+print(json.dumps(details, indent=4, ensure_ascii=False))

examples/joyai_image/model_inference/JoyAI-Image-Edit.py

@@ -0,0 +1,39 @@
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+import torch
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+# Download dataset
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="joyai_image/JoyAI-Image-Edit/*"
+)
+
+pipe = JoyAIImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model*.safetensors"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="vae/Wan2.1_VAE.pth"),
+    ],
+    processor_config=ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/"),
+)
+
+# Use first sample from dataset
+dataset_base_path = "data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit"
+prompt = "将裙子改为粉色"
+edit_image = Image.open(f"{dataset_base_path}/edit/image1.jpg").convert("RGB")
+
+output = pipe(
+    prompt=prompt,
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=1,
+    num_inference_steps=30,
+    cfg_scale=5.0,
+)
+
+output.save("output_joyai_edit.png")

examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py

@@ -0,0 +1,51 @@
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+import torch
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+# Download dataset
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="joyai_image/JoyAI-Image-Edit/*"
+)
+
+vram_config = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cpu",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+
+pipe = JoyAIImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth", **vram_config),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model*.safetensors", **vram_config),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="vae/Wan2.1_VAE.pth", **vram_config),
+    ],
+    processor_config=ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+# Use first sample from dataset
+dataset_base_path = "data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit"
+prompt = "将裙子改为粉色"
+edit_image = Image.open(f"{dataset_base_path}/edit/image1.jpg").convert("RGB")
+
+output = pipe(
+    prompt=prompt,
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=0,
+    num_inference_steps=30,
+    cfg_scale=5.0,
+)
+
+output.save("output_joyai_edit_low_vram.png")

examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh

@@ -0,0 +1,35 @@
+# Dataset: data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/
+# Download: modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "joyai_image/JoyAI-Image-Edit/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/joyai_image/model_training/train.py \
+    --dataset_base_path "./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit" \
+    --dataset_metadata_path "./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/metadata.csv" \
+    --max_pixels 1048576 \
+    --dataset_repeat 1 \
+    --model_id_with_origin_paths "jd-opensource/JoyAI-Image-Edit:JoyAI-Image-Und/model*.safetensors,jd-opensource/JoyAI-Image-Edit:vae/Wan2.1_VAE.pth" \
+    --learning_rate 1e-5 \
+    --num_epochs 2 \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --output_path "./models/train/JoyAI-Image-Edit-full-cache" \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --data_file_keys "image,edit_image" \
+    --extra_inputs "edit_image" \
+    --task "sft:data_process"
+
+accelerate launch --config_file examples/joyai_image/model_training/full/accelerate_config_zero3.yaml \
+    examples/joyai_image/model_training/train.py \
+    --dataset_base_path "./models/train/JoyAI-Image-Edit-full-cache" \
+    --max_pixels 1048576 \
+    --dataset_repeat 50 \
+    --model_id_with_origin_paths "jd-opensource/JoyAI-Image-Edit:transformer/transformer.pth" \
+    --learning_rate 1e-5 \
+    --num_epochs 2 \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --output_path "./models/train/JoyAI-Image-Edit-full" \
+    --trainable_models "dit" \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --data_file_keys "image,edit_image" \
+    --extra_inputs "edit_image" \
+    --task "sft:train"

examples/joyai_image/model_training/full/accelerate_config_zero3.yaml

@@ -0,0 +1,23 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh

@@ -0,0 +1,39 @@
+# Dataset: data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/
+# Download: modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "joyai_image/JoyAI-Image-Edit/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/joyai_image/model_training/train.py \
+    --dataset_base_path "./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit" \
+    --dataset_metadata_path "./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/metadata.csv" \
+    --max_pixels 1048576 \
+    --dataset_repeat 1 \
+    --model_id_with_origin_paths "jd-opensource/JoyAI-Image-Edit:JoyAI-Image-Und/model*.safetensors,jd-opensource/JoyAI-Image-Edit:vae/Wan2.1_VAE.pth" \
+    --learning_rate 1e-4 \
+    --num_epochs 5 \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --output_path "./models/train/JoyAI-Image-Edit-split-cache" \
+    --lora_base_model "dit" \
+    --lora_target_modules "img_attn_qkv,txt_attn_qkv,img_attn_proj,txt_attn_proj" \
+    --lora_rank 32 \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --data_file_keys "image,edit_image" \
+    --extra_inputs "edit_image" \
+    --task "sft:data_process"
+
+accelerate launch examples/joyai_image/model_training/train.py \
+    --dataset_base_path "./models/train/JoyAI-Image-Edit-split-cache" \
+    --max_pixels 1048576 \
+    --dataset_repeat 50 \
+    --model_id_with_origin_paths "jd-opensource/JoyAI-Image-Edit:transformer/transformer.pth" \
+    --learning_rate 1e-4 \
+    --num_epochs 5 \
+    --remove_prefix_in_ckpt "pipe.dit." \
+    --output_path "./models/train/JoyAI-Image-Edit-lora" \
+    --lora_base_model "dit" \
+    --lora_target_modules "img_attn_qkv,txt_attn_qkv,img_attn_proj,txt_attn_proj" \
+    --lora_rank 32 \
+    --use_gradient_checkpointing \
+    --find_unused_parameters \
+    --data_file_keys "image,edit_image" \
+    --extra_inputs "edit_image" \
+    --task "sft:train"

examples/joyai_image/model_training/special/split_training/JoyAI-Image-Edit.sh

@@ -0,0 +1,42 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "joyai_image/JoyAI-Image-Edit/*" --local_dir ./data/diffsynth_example_dataset
+
+# Stage 1: cache deterministic preprocessing outputs.
+accelerate launch examples/joyai_image/model_training/train.py \
+  --dataset_base_path ./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit \
+  --dataset_metadata_path ./data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 1 \
+  --model_id_with_origin_paths 'jd-opensource/JoyAI-Image-Edit:transformer/transformer.pth,jd-opensource/JoyAI-Image-Edit:JoyAI-Image-Und/model*.safetensors,jd-opensource/JoyAI-Image-Edit:vae/Wan2.1_VAE.pth' \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/JoyAI-Image-Edit-split-cache \
+  --lora_base_model dit \
+  --lora_target_modules img_attn_qkv,txt_attn_qkv,img_attn_proj,txt_attn_proj \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --data_file_keys image,edit_image \
+  --extra_inputs edit_image \
+  --task sft:data_process \
+  --offload_models jd-opensource/JoyAI-Image-Edit:transformer/transformer.pth
+
+# Stage 2: train LoRA from the cached dataset.
+accelerate launch examples/joyai_image/model_training/train.py \
+  --dataset_base_path ./models/train/JoyAI-Image-Edit-split-cache \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths 'jd-opensource/JoyAI-Image-Edit:transformer/transformer.pth,jd-opensource/JoyAI-Image-Edit:JoyAI-Image-Und/model*.safetensors,jd-opensource/JoyAI-Image-Edit:vae/Wan2.1_VAE.pth' \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/JoyAI-Image-Edit-split \
+  --lora_base_model dit \
+  --lora_target_modules img_attn_qkv,txt_attn_qkv,img_attn_proj,txt_attn_proj \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --data_file_keys image,edit_image \
+  --extra_inputs edit_image \
+  --task sft:train \
+  --offload_models 'jd-opensource/JoyAI-Image-Edit:JoyAI-Image-Und/model*.safetensors,jd-opensource/JoyAI-Image-Edit:vae/Wan2.1_VAE.pth'

examples/joyai_image/model_training/special/split_training/validate.py

@@ -0,0 +1,30 @@
+import torch
+from PIL import Image
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+
+pipe = JoyAIImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model*.safetensors"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="vae/Wan2.1_VAE.pth"),
+    ],
+    processor_config=ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/"),
+)
+
+pipe.load_lora(pipe.dit, './models/train/JoyAI-Image-Edit-split/epoch-4.safetensors')
+
+prompt = "将裙子改为粉色"
+edit_image = Image.open("data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/edit/image1.jpg").convert("RGB")
+
+image = pipe(
+    prompt=prompt,
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=0,
+    num_inference_steps=30,
+    cfg_scale=5.0,
+)
+image.save('split_training_JoyAI-Image-Edit.jpg')

examples/joyai_image/model_training/train.py

@@ -0,0 +1,150 @@
+import torch, os, argparse, accelerate
+from diffsynth.core import UnifiedDataset
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+from diffsynth.diffusion import *
+from diffsynth.core.data.operators import *
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class JoyAIImageTrainingModule(DiffusionTrainingModule):
+    def __init__(
+        self,
+        model_paths=None, model_id_with_origin_paths=None,
+        processor_path=None,
+        trainable_models=None,
+        lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None,
+        preset_lora_path=None, preset_lora_model=None,
+        use_gradient_checkpointing=True,
+        use_gradient_checkpointing_offload=False,
+        extra_inputs=None,
+        fp8_models=None,
+        offload_models=None,
+        quant_options=None,
+        resume_from_checkpoint=None, remove_prefix_in_ckpt=None,
+        device="cpu",
+        task="sft",
+    ):
+        super().__init__()
+        # Load models
+        model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, quant_options=quant_options, device=device)
+        processor_config = ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/") if processor_path is None else ModelConfig(processor_path)
+        self.pipe = JoyAIImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, processor_config=processor_config)
+        self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model)
+        self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt)
+
+        # Training mode
+        self.switch_pipe_to_training_mode(
+            self.pipe, trainable_models,
+            lora_base_model, lora_target_modules, lora_rank, lora_checkpoint,
+            preset_lora_path, preset_lora_model,
+            task=task,
+        )
+
+        # Other configs
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
+        self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else []
+        self.fp8_models = fp8_models
+        self.task = task
+        self.task_to_loss = {
+            "sft:data_process": lambda pipe, *args: args,
+            "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+            "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+        }
+
+    def get_pipeline_inputs(self, data):
+        inputs_posi = {"prompt": data["prompt"]}
+        inputs_nega = {"negative_prompt": ""}
+        inputs_shared = {
+            # Assume you are using this pipeline for inference,
+            # please fill in the input parameters.
+            "input_image": data["image"],
+            "height": data["image"].size[1],
+            "width": data["image"].size[0],
+            # Please do not modify the following parameters
+            # unless you clearly know what this will cause.
+            "cfg_scale": 1,
+            "rand_device": self.pipe.device,
+            "use_gradient_checkpointing": self.use_gradient_checkpointing,
+            "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload,
+        }
+        inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared)
+        return inputs_shared, inputs_posi, inputs_nega
+
+    def forward(self, data, inputs=None):
+        if inputs is None: inputs = self.get_pipeline_inputs(data)
+        inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype)
+        for unit in self.pipe.units:
+            inputs = self.pipe.unit_runner(unit, self.pipe, *inputs)
+        loss = self.task_to_loss[self.task](self.pipe, *inputs)
+        return loss
+
+
+def joyai_image_parser():
+    parser = argparse.ArgumentParser(description="JoyAI-Image training.")
+    parser = add_general_config(parser)
+    parser = add_image_size_config(parser)
+    parser.add_argument("--processor_path", type=str, default=None, help="Path to the processor.")
+    parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
+    return parser
+
+
+if __name__ == "__main__":
+    parser = joyai_image_parser()
+    args = parser.parse_args()
+    accelerator = accelerate.Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)],
+    )
+    dataset = UnifiedDataset(
+        base_path=args.dataset_base_path,
+        metadata_path=args.dataset_metadata_path,
+        repeat=args.dataset_repeat,
+        data_file_keys=args.data_file_keys.split(","),
+        main_data_operator=UnifiedDataset.default_image_operator(
+            base_path=args.dataset_base_path,
+            max_pixels=args.max_pixels,
+            height=args.height,
+            width=args.width,
+            height_division_factor=16,
+            width_division_factor=16,
+        ),
+    )
+    model = JoyAIImageTrainingModule(
+        model_paths=args.model_paths,
+        model_id_with_origin_paths=args.model_id_with_origin_paths,
+        processor_path=args.processor_path,
+        trainable_models=args.trainable_models,
+        lora_base_model=args.lora_base_model,
+        lora_target_modules=args.lora_target_modules,
+        lora_rank=args.lora_rank,
+        lora_checkpoint=args.lora_checkpoint,
+        preset_lora_path=args.preset_lora_path,
+        preset_lora_model=args.preset_lora_model,
+        use_gradient_checkpointing=args.use_gradient_checkpointing,
+        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
+        extra_inputs=args.extra_inputs,
+        fp8_models=args.fp8_models,
+        offload_models=args.offload_models,
+        quant_options=args.quant_options,
+        resume_from_checkpoint=args.resume_from_checkpoint,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        task=args.task,
+        device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device,
+    )
+    model_logger = ModelLogger(
+        args.output_path,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        enable_tensorboard_log=args.enable_tensorboard_log,
+        enable_swanlab_log=args.enable_swanlab_log,
+        swanlab_project=args.swanlab_project,
+        enable_wandb_log=args.enable_wandb_log,
+        wandb_project=args.wandb_project,
+        enable_csv_log=args.enable_csv_log,
+    )
+    launcher_map = {
+        "sft:data_process": launch_data_process_task,
+        "sft": launch_training_task,
+        "sft:train": launch_training_task,
+    }
+    launcher_map[args.task](accelerator, dataset, model, model_logger, args=args)

examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py

@@ -0,0 +1,32 @@
+import torch
+from PIL import Image
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+from diffsynth import load_state_dict
+
+pipe = JoyAIImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model*.safetensors"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="vae/Wan2.1_VAE.pth"),
+    ],
+    processor_config=ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/"),
+)
+
+state_dict = load_state_dict("models/train/JoyAI-Image-Edit_full/epoch-1.safetensors")
+pipe.dit.load_state_dict(state_dict)
+
+prompt = "将裙子改为粉色"
+edit_image = Image.open("data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/edit/image1.jpg").convert("RGB")
+
+image = pipe(
+    prompt=prompt,
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=0,
+    num_inference_steps=50,
+    cfg_scale=5.0,
+)
+image.save("image_full.jpg")

examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py

@@ -0,0 +1,30 @@
+import torch
+from PIL import Image
+from diffsynth.pipelines.joyai_image import JoyAIImagePipeline, ModelConfig
+
+pipe = JoyAIImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="transformer/transformer.pth"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/model*.safetensors"),
+        ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="vae/Wan2.1_VAE.pth"),
+    ],
+    processor_config=ModelConfig(model_id="jd-opensource/JoyAI-Image-Edit", origin_file_pattern="JoyAI-Image-Und/"),
+)
+
+pipe.load_lora(pipe.dit, "models/train/JoyAI-Image-Edit-lora/epoch-4.safetensors")
+
+prompt = "将裙子改为粉色"
+edit_image = Image.open("data/diffsynth_example_dataset/joyai_image/JoyAI-Image-Edit/edit/image1.jpg").convert("RGB")
+
+image = pipe(
+    prompt=prompt,
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=0,
+    num_inference_steps=30,
+    cfg_scale=5.0,
+)
+image.save("image_lora.jpg")

examples/krea2/model_inference/Krea-2-Raw.py

@@ -0,0 +1,17 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+prompt = "A cat standing on a stone."
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/krea2/model_inference/Krea-2-Turbo.py

@@ -0,0 +1,22 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Turbo", origin_file_pattern="turbo.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+prompt = "Portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles."
+image = pipe(
+    prompt, seed=0,
+    height=2048, width=2048,
+    # The following parameters are fixed.
+    num_inference_steps=8, cfg_scale=1, mu=1.15,
+)
+image.save("image.jpg")

examples/krea2/model_inference_low_vram/Krea-2-Raw.py

@@ -0,0 +1,27 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+vram_config = {
+    "offload_dtype": "disk",
+    "offload_device": "disk",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 1,
+)
+prompt = "A cat standing on a stone."
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/krea2/model_inference_low_vram/Krea-2-Turbo.py

@@ -0,0 +1,32 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+vram_config = {
+    "offload_dtype": "disk",
+    "offload_device": "disk",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Turbo", origin_file_pattern="turbo.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 1,
+)
+prompt = "Portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles."
+image = pipe(
+    prompt, seed=0,
+    height=2048, width=2048,
+    # The following parameters are fixed.
+    num_inference_steps=8, cfg_scale=1, mu=1.15,
+)
+image.save("image.jpg")

examples/krea2/model_training/full/Krea-2-Raw.sh

@@ -0,0 +1,17 @@
+# Please run `accelerate config` to configure GPU, DeepSpeed, etc.
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "krea2/Krea-2-Raw/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/krea2/Krea-2-Raw \
+  --dataset_metadata_path data/diffsynth_example_dataset/krea2/Krea-2-Raw/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "krea/Krea-2-Raw:raw.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
+  --tokenizer_path "Qwen/Qwen3-VL-4B-Instruct:" \
+  --learning_rate 1e-5 \
+  --num_epochs 2 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Krea-2-Raw_full" \
+  --trainable_models "dit" \
+  --use_gradient_checkpointing \
+  --find_unused_parameters

examples/krea2/model_training/full/Krea-2-Turbo.sh

@@ -0,0 +1,17 @@
+# Please run `accelerate config` to configure GPU, DeepSpeed, etc.
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "krea2/Krea-2-Turbo/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/krea2/Krea-2-Turbo \
+  --dataset_metadata_path data/diffsynth_example_dataset/krea2/Krea-2-Turbo/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "krea/Krea-2-Turbo:turbo.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
+  --tokenizer_path "Qwen/Qwen3-VL-4B-Instruct:" \
+  --learning_rate 1e-5 \
+  --num_epochs 2 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Krea-2-Turbo_full" \
+  --trainable_models "dit" \
+  --use_gradient_checkpointing \
+  --find_unused_parameters

examples/krea2/model_training/lora/Krea-2-Raw.sh

@@ -0,0 +1,19 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "krea2/Krea-2-Raw/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/krea2/Krea-2-Raw \
+  --dataset_metadata_path data/diffsynth_example_dataset/krea2/Krea-2-Raw/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "krea/Krea-2-Raw:raw.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
+  --tokenizer_path "Qwen/Qwen3-VL-4B-Instruct:" \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Krea-2-Raw_lora" \
+  --lora_base_model "dit" \
+  --lora_target_modules "wq,wk,wv,gate,wo,gate,up,down,first,tmlp.0,tmlp.2,projector,txtmlp.1,txtmlp.3,last.linear,tproj.1" \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --align_to_opensource_format

examples/krea2/model_training/lora/Krea-2-Turbo.sh

@@ -0,0 +1,21 @@
+# We do not recommend training LoRA on Krea-2-Turbo.
+
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "krea2/Krea-2-Turbo/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/krea2/Krea-2-Turbo \
+  --dataset_metadata_path data/diffsynth_example_dataset/krea2/Krea-2-Turbo/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "krea/Krea-2-Turbo:turbo.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
+  --tokenizer_path "Qwen/Qwen3-VL-4B-Instruct:" \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Krea-2-Turbo_lora" \
+  --lora_base_model "dit" \
+  --lora_target_modules "wq,wk,wv,gate,wo,gate,up,down,first,tmlp.0,tmlp.2,projector,txtmlp.1,txtmlp.3,last.linear,tproj.1" \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --align_to_opensource_format

examples/krea2/model_training/special/split_training/Krea-2-Raw.sh

@@ -0,0 +1,42 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "krea2/Krea-2-Raw/*" --local_dir ./data/diffsynth_example_dataset
+
+# Stage 1: cache deterministic preprocessing outputs.
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/krea2/Krea-2-Raw \
+  --dataset_metadata_path data/diffsynth_example_dataset/krea2/Krea-2-Raw/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 1 \
+  --model_id_with_origin_paths 'krea/Krea-2-Raw:raw.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors' \
+  --tokenizer_path Qwen/Qwen3-VL-4B-Instruct: \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/Krea-2-Raw_split_cache \
+  --lora_base_model dit \
+  --lora_target_modules wq,wk,wv,gate,wo,gate,up,down,first,tmlp.0,tmlp.2,projector,txtmlp.1,txtmlp.3,last.linear,tproj.1 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --align_to_opensource_format \
+  --offload_models krea/Krea-2-Raw:raw.safetensors \
+  --task sft:data_process
+
+# Stage 2: train LoRA from the cached dataset.
+accelerate launch examples/krea2/model_training/train.py \
+  --dataset_base_path ./models/train/Krea-2-Raw_split_cache \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths 'krea/Krea-2-Raw:raw.safetensors,Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors' \
+  --tokenizer_path Qwen/Qwen3-VL-4B-Instruct: \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/Krea-2-Raw_split \
+  --lora_base_model dit \
+  --lora_target_modules wq,wk,wv,gate,wo,gate,up,down,first,tmlp.0,tmlp.2,projector,txtmlp.1,txtmlp.3,last.linear,tproj.1 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --align_to_opensource_format \
+  --offload_models 'Qwen/Qwen3-VL-4B-Instruct:*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors' \
+  --task sft:train

examples/krea2/model_training/special/split_training/validate.py

@@ -0,0 +1,19 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        # For LoRA models trained on Krea-2-Raw, we recommend using them on Krea-2-Turbo.
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+pipe.load_lora(pipe.dit, './models/train/Krea-2-Raw_split/epoch-4.safetensors')
+prompt = "A dog"
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save('split_training_Krea-2-Raw.jpg')

examples/krea2/model_training/train.py

@@ -0,0 +1,151 @@
+import torch, os, argparse, accelerate
+from diffsynth.core import UnifiedDataset
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+from diffsynth.diffusion import *
+from diffsynth.utils.lora.krea2 import Krea2LoRAConverter
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class Krea2ImageTrainingModule(DiffusionTrainingModule):
+    def __init__(
+        self,
+        model_paths=None, model_id_with_origin_paths=None,
+        tokenizer_path=None,
+        trainable_models=None,
+        lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None,
+        preset_lora_path=None, preset_lora_model=None,
+        use_gradient_checkpointing=True,
+        use_gradient_checkpointing_offload=False,
+        extra_inputs=None,
+        fp8_models=None,
+        offload_models=None,
+        quant_options=None,
+        template_model_id_or_path=None,
+        resume_from_checkpoint=None, remove_prefix_in_ckpt=None,
+        enable_lora_hot_loading=False,
+        device="cpu",
+        task="sft",
+    ):
+        super().__init__()
+        model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, quant_options=quant_options, device=device)
+        tokenizer_config = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""))
+        self.pipe = Krea2Pipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config)
+        self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, use_gradient_checkpointing, use_gradient_checkpointing_offload)
+        self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model, remove_unnecessary_params=True)
+        self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt)
+        if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit)
+
+        self.switch_pipe_to_training_mode(
+            self.pipe, trainable_models,
+            lora_base_model, lora_target_modules, lora_rank, lora_checkpoint,
+            preset_lora_path, preset_lora_model,
+            task=task,
+        )
+
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
+        self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else []
+        self.fp8_models = fp8_models
+        self.task = task
+        self.task_to_loss = {
+            "sft:data_process": lambda pipe, *args: args,
+            "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+            "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
+        }
+
+    def get_pipeline_inputs(self, data):
+        inputs_posi = {"prompt": data["prompt"]}
+        inputs_nega = {"negative_prompt": ""}
+        inputs_shared = {
+            "input_image": data["image"],
+            "height": data["image"].size[1],
+            "width": data["image"].size[0],
+            "cfg_scale": 1,
+            "rand_device": self.pipe.device,
+            "use_gradient_checkpointing": self.use_gradient_checkpointing,
+            "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload,
+        }
+        inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared)
+        return inputs_shared, inputs_posi, inputs_nega
+
+    def forward(self, data, inputs=None):
+        if inputs is None: inputs = self.get_pipeline_inputs(data)
+        inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype)
+        for unit in self.pipe.units:
+            inputs = self.pipe.unit_runner(unit, self.pipe, *inputs)
+        loss = self.task_to_loss[self.task](self.pipe, *inputs)
+        return loss
+
+
+def krea2_parser():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser = add_general_config(parser)
+    parser = add_image_size_config(parser)
+    parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.")
+    parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
+    parser.add_argument("--align_to_opensource_format", default=False, action="store_true", help="Whether to align the lora format to opensource format. Example: https://modelscope.cn/models/krea/Krea-2-LoRA-softwatercolor")
+    return parser
+
+
+if __name__ == "__main__":
+    parser = krea2_parser()
+    args = parser.parse_args()
+    accelerator = accelerate.Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)],
+    )
+    dataset = UnifiedDataset(
+        base_path=args.dataset_base_path,
+        metadata_path=args.dataset_metadata_path,
+        repeat=args.dataset_repeat,
+        data_file_keys=args.data_file_keys.split(","),
+        main_data_operator=UnifiedDataset.default_image_operator(
+            base_path=args.dataset_base_path,
+            max_pixels=args.max_pixels,
+            height=args.height,
+            width=args.width,
+            height_division_factor=16,
+            width_division_factor=16,
+        )
+    )
+    model = Krea2ImageTrainingModule(
+        model_paths=args.model_paths,
+        model_id_with_origin_paths=args.model_id_with_origin_paths,
+        tokenizer_path=args.tokenizer_path,
+        trainable_models=args.trainable_models,
+        lora_base_model=args.lora_base_model,
+        lora_target_modules=args.lora_target_modules,
+        lora_rank=args.lora_rank,
+        lora_checkpoint=args.lora_checkpoint,
+        preset_lora_path=args.preset_lora_path,
+        preset_lora_model=args.preset_lora_model,
+        use_gradient_checkpointing=args.use_gradient_checkpointing,
+        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
+        extra_inputs=args.extra_inputs,
+        fp8_models=args.fp8_models,
+        offload_models=args.offload_models,
+        quant_options=args.quant_options,
+        template_model_id_or_path=args.template_model_id_or_path,
+        resume_from_checkpoint=args.resume_from_checkpoint,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        enable_lora_hot_loading=args.enable_lora_hot_loading,
+        task=args.task,
+        device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device,
+    )
+    model_logger = ModelLogger(
+        args.output_path,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        state_dict_converter=Krea2LoRAConverter.align_to_opensource_format if args.align_to_opensource_format else lambda x:x,
+        enable_tensorboard_log=args.enable_tensorboard_log,
+        enable_swanlab_log=args.enable_swanlab_log,
+        swanlab_project=args.swanlab_project,
+        enable_wandb_log=args.enable_wandb_log,
+        wandb_project=args.wandb_project,
+        enable_csv_log=args.enable_csv_log,
+    )
+    launcher_map = {
+        "sft:data_process": launch_data_process_task,
+        "sft": launch_training_task,
+        "sft:train": launch_training_task,
+    }
+    launcher_map[args.task](accelerator, dataset, model, model_logger, args=args)

examples/krea2/model_training/validate_full/Krea-2-Raw.py

@@ -0,0 +1,19 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+from diffsynth import load_state_dict
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+pipe.dit.load_state_dict(load_state_dict("models/train/Krea-2-Raw_full/epoch-1.safetensors"))
+prompt = "A dog"
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/krea2/model_training/validate_full/Krea-2-Turbo.py

@@ -0,0 +1,19 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+from diffsynth import load_state_dict
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Turbo", origin_file_pattern="turbo.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+pipe.dit.load_state_dict(load_state_dict("models/train/Krea-2-Turbo_full/epoch-1.safetensors"))
+prompt = "A dog"
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/krea2/model_training/validate_lora/Krea-2-Raw.py

@@ -0,0 +1,19 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        # For LoRA models trained on Krea-2-Raw, we recommend using them on Krea-2-Turbo.
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+pipe.load_lora(pipe.dit, "models/train/Krea-2-Raw_lora/epoch-4.safetensors")
+prompt = "A dog"
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/krea2/model_training/validate_lora/Krea-2-Turbo.py

@@ -0,0 +1,18 @@
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Turbo", origin_file_pattern="turbo.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+pipe.load_lora(pipe.dit, "models/train/Krea-2-Turbo_lora/epoch-4.safetensors")
+prompt = "A dog"
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")

examples/lingbot_video/README.md

@@ -0,0 +1,3 @@
+English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LingBot-Video.html
+
+中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/LingBot-Video.html

examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2i.py

@@ -0,0 +1,32 @@
+import json
+import torch
+from diffsynth.pipelines.lingbot_video import LingBotVideoPipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+pipe = LingBotVideoPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="lingbot_video/lingbot-video-dense-1.3b_t2i/*",
+)
+with open("data/diffsynth_example_dataset/lingbot_video/lingbot-video-dense-1.3b_t2i/t2i_example.json", "r", encoding="utf-8") as f:
+    caption = json.load(f)
+
+frames = pipe(
+    prompt=caption,
+    negative_prompt=pipe.default_negative_prompt_image,
+    height=480, width=832, num_frames=1,
+    num_inference_steps=40, cfg_scale=3.0,
+    seed=0,
+)
+frames[0].save("image_lingbot-video-dense-1.3b_t2i.png")

examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_t2v.py

@@ -0,0 +1,68 @@
+import torch
+import json
+from diffsynth.utils.data import save_video, VideoData
+from diffsynth.pipelines.lingbot_video import LingBotVideoPipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+pipe = LingBotVideoPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+
+# --- Text-to-video -------------------------------------------------------------------
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="lingbot_video/lingbot-video-dense-1.3b_t2v/*",
+)
+# LingBot-Video is trained on structured-JSON captions, not free-form prose. This example
+# runs on a released in-distribution caption; see the bottom for turning a brief idea into
+# such a caption with the two-stage prompt rewriter.
+with open("data/diffsynth_example_dataset/lingbot_video/lingbot-video-dense-1.3b_t2v/t2v_example_1.json", "r", encoding="utf-8") as f:
+    caption = json.load(f)
+
+video = pipe(
+    prompt=caption,
+    negative_prompt=pipe.default_negative_prompt,
+    height=480, width=832, num_frames=81,
+    num_inference_steps=40, cfg_scale=3.0,
+    seed=0,
+)
+save_video(video, "video_lingbot-video-dense-1.3b_t2v.mp4", fps=15, quality=10)
+
+# --- Video-to-video ------------------------------------------------------------------
+# denoising_strength < 1 keeps part of the input structure.
+input_video = VideoData("video_lingbot-video-dense-1.3b_t2v.mp4", height=480, width=832)
+video = pipe(
+    prompt=caption,
+    negative_prompt=pipe.default_negative_prompt,
+    input_video=input_video, denoising_strength=0.7,
+    height=480, width=832, num_frames=81,
+    num_inference_steps=40, cfg_scale=3.0,
+    seed=1,
+)
+save_video(video, "video_lingbot-video-dense-1.3b_v2v.mp4", fps=15, quality=10)
+
+# --- Optional: rewrite a brief idea into a structured caption ------------------------
+# The two-stage rewriter (model_training/scripts/prompt_rewriter.py) is a separate VLM +
+# stage-2 LoRA adapter (NOT the DiT) and is not downloaded automatically. Fetch both
+# weights and point the env vars at them:
+#     modelscope download --model Qwen/Qwen3.6-27B --local_dir ./models/Qwen/Qwen3.6-27B
+#     modelscope download --model Robbyant/lingbot-video-rewriter-lora --local_dir ./models/Robbyant/lingbot-video-rewriter-lora
+#     export REWRITER_BASE_MODEL=./models/Qwen/Qwen3.6-27B
+#     export REWRITER_ADAPTER=./models/Robbyant/lingbot-video-rewriter-lora
+#
+# from examples.lingbot_video.model_training.scripts.prompt_rewriter import rewrite_prompt
+# caption = rewrite_prompt(
+#     "A playful puppy runs across a lush green meadow, chasing a red ball. "
+#     "Dynamic side-tracking camera.",
+#     mode="t2v", duration=5,
+# )
+# video = pipe(prompt=caption, height=480, width=832, num_frames=81, cfg_scale=3.0, seed=0)
+# save_video(video, "video_lingbot-video-dense-1.3b_rewrite.mp4", fps=15, quality=10)

examples/lingbot_video/model_inference/lingbot-video-dense-1.3b_ti2v.py

@@ -0,0 +1,39 @@
+import os
+import json
+import torch
+from PIL import Image
+from diffsynth.utils.data import save_video
+from diffsynth.pipelines.lingbot_video import LingBotVideoPipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+pipe = LingBotVideoPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Robbyant/lingbot-video-dense-1.3b", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+
+# The condition first frame and its paired caption ship in the example dataset.
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="lingbot_video/lingbot-video-dense-1.3b_ti2v/*",
+)
+base = "data/diffsynth_example_dataset/lingbot_video/lingbot-video-dense-1.3b_ti2v"
+with open(os.path.join(base, "ti2v_example.json"), "r", encoding="utf-8") as f:
+    caption = json.load(f)
+input_image = Image.open(os.path.join(base, "ti2v_first_frame.png")).convert("RGB")
+
+video = pipe(
+    prompt=caption,
+    negative_prompt=pipe.default_negative_prompt,
+    input_image=input_image,
+    height=480, width=832, num_frames=81,
+    num_inference_steps=40, cfg_scale=3.0,
+    seed=0,
+)
+save_video(video, "video_lingbot-video-dense-1.3b_ti2v.mp4", fps=15, quality=10)

examples/lingbot_video/model_inference/lingbot-video-moe-30b-a3b_t2i.py

@@ -0,0 +1,32 @@
+import json
+import torch
+from diffsynth.pipelines.lingbot_video import LingBotVideoPipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+pipe = LingBotVideoPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Robbyant/lingbot-video-moe-30b-a3b", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors"),
+        ModelConfig(model_id="Robbyant/lingbot-video-moe-30b-a3b", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+)
+
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/diffsynth_example_dataset",
+    local_dir="data/diffsynth_example_dataset",
+    allow_file_pattern="lingbot_video/lingbot-video-moe-30b-a3b_t2i/*",
+)
+with open("data/diffsynth_example_dataset/lingbot_video/lingbot-video-moe-30b-a3b_t2i/t2i_example.json", "r", encoding="utf-8") as f:
+    caption = json.load(f)
+
+frames = pipe(
+    prompt=caption,
+    negative_prompt=pipe.default_negative_prompt_image,
+    height=480, width=832, num_frames=1,
+    num_inference_steps=40, cfg_scale=3.0,
+    seed=0,
+)
+frames[0].save("image_lingbot-video-moe-30b-a3b_t2i.png")