Upload folder using huggingface_hub (part 3)

examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Edit.py

@@ -0,0 +1,32 @@
+import torch
+from PIL import Image
+from diffsynth.pipelines.boogu_image import BooguImagePipeline, ModelConfig
+
+pipe = BooguImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Edit", origin_file_pattern="transformer/*.safetensors"),
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Edit", origin_file_pattern="mllm/*.safetensors"),
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Edit", origin_file_pattern="vae/*.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Boogu/Boogu-Image-0.1-Edit", origin_file_pattern="mllm/"),
+)
+
+pipe.load_lora(pipe.dit, "models/train/Boogu-Image-0.1-Edit_lora/epoch-4.safetensors")
+
+prompt = "将裙子改为粉色"
+edit_image = Image.open("data/diffsynth_example_dataset/boogu_image/Boogu-Image-0.1-Edit/edit/image1.jpg").convert("RGB")
+
+output = pipe(
+    prompt=prompt,
+    negative_prompt="",
+    edit_image=edit_image,
+    height=1024,
+    width=1024,
+    seed=42,
+    rand_device="cuda",
+    num_inference_steps=50,
+    cfg_scale=1.0,
+)
+output.save("image_Boogu-Image-0.1-Edit_lora.jpg")

examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Turbo.py

@@ -0,0 +1,30 @@
+import torch
+from diffsynth.pipelines.boogu_image import BooguImagePipeline, ModelConfig
+
+pipe = BooguImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Turbo", origin_file_pattern="transformer/*.safetensors"),
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Turbo", origin_file_pattern="mllm/*.safetensors"),
+        ModelConfig(model_id="Boogu/Boogu-Image-0.1-Turbo", origin_file_pattern="vae/*.safetensors"),
+    ],
+    processor_config=ModelConfig(model_id="Boogu/Boogu-Image-0.1-Turbo", origin_file_pattern="mllm/"),
+)
+
+pipe.load_lora(pipe.dit, "models/train/Boogu-Image-0.1-Turbo_lora/epoch-4.safetensors")
+
+prompt = "dog,white and brown dog, sitting on wall, under pink flowers"
+
+output = pipe(
+    prompt=prompt,
+    negative_prompt="",
+    height=1024,
+    width=1024,
+    seed=42,
+    rand_device="cuda",
+    num_inference_steps=4,
+    cfg_scale=1.0,
+    sigmas=[0.999, 0.748, 0.5, 0.25],
+)
+output.save("image_Boogu-Image-0.1-Turbo_lora.jpg")

examples/dev_tools/unit_test.py

@@ -0,0 +1,121 @@
+import os, shutil, multiprocessing, time
+NUM_GPUS = 7
+
+
+def script_is_processed(output_path, script):
+    return os.path.exists(os.path.join(output_path, script)) and "log.txt" in os.listdir(os.path.join(output_path, script))
+
+
+def filter_unprocessed_tasks(script_path):
+    tasks = []
+    output_path = os.path.join("data", script_path)
+    for script in sorted(os.listdir(script_path)):
+        if not script.endswith(".sh") and not script.endswith(".py"):
+            continue
+        if script_is_processed(output_path, script):
+            continue
+        tasks.append(script)
+    return tasks
+
+
+def run_inference(script_path):
+    tasks = filter_unprocessed_tasks(script_path)
+    output_path = os.path.join("data", script_path)
+    for script in tasks:
+        source_path = os.path.join(script_path, script)
+        target_path = os.path.join(output_path, script)
+        os.makedirs(target_path, exist_ok=True)
+        cmd = f"python {source_path} > {target_path}/log.txt 2>&1"
+        print(cmd, flush=True)
+        os.system(cmd)
+        for file_name in os.listdir("./"):
+            if file_name.endswith(".jpg") or file_name.endswith(".png") or file_name.endswith(".mp4"):
+                shutil.move(file_name, os.path.join(target_path, file_name))
+
+
+def run_tasks_on_single_GPU(script_path, tasks, gpu_id, num_gpu):
+    output_path = os.path.join("data", script_path)
+    for script_id, script in enumerate(tasks):
+        if script_id % num_gpu != gpu_id:
+            continue
+        source_path = os.path.join(script_path, script)
+        target_path = os.path.join(output_path, script)
+        os.makedirs(target_path, exist_ok=True)
+        if script.endswith(".sh"):
+            cmd = f"CUDA_VISIBLE_DEVICES={gpu_id} bash {source_path} > {target_path}/log.txt 2>&1"
+        elif script.endswith(".py"):
+            cmd = f"CUDA_VISIBLE_DEVICES={gpu_id} python {source_path} > {target_path}/log.txt 2>&1"
+        print(cmd, flush=True)
+        os.system(cmd)
+
+
+def run_train_multi_GPU(script_path):
+    tasks = filter_unprocessed_tasks(script_path)
+    output_path = os.path.join("data", script_path)
+    for script in tasks:
+        source_path = os.path.join(script_path, script)
+        target_path = os.path.join(output_path, script)
+        os.makedirs(target_path, exist_ok=True)
+        cmd = f"bash {source_path} > {target_path}/log.txt 2>&1"
+        print(cmd, flush=True)
+        os.system(cmd)
+        time.sleep(1)
+        
+
+def run_train_single_GPU(script_path):
+    tasks = filter_unprocessed_tasks(script_path)
+    processes = [multiprocessing.Process(target=run_tasks_on_single_GPU, args=(script_path, tasks, i, NUM_GPUS)) for i in range(NUM_GPUS)]
+    for p in processes:
+        p.start()
+    for p in processes:
+        p.join()
+
+
+def move_files(prefix, target_folder):
+    os.makedirs(target_folder, exist_ok=True)
+    os.system(f"cp -r {prefix}* {target_folder}")
+    os.system(f"rm -rf {prefix}*")
+
+
+def test_qwen_image():
+    run_inference("examples/qwen_image/model_inference")
+    run_inference("examples/qwen_image/model_inference_low_vram")
+    run_train_multi_GPU("examples/qwen_image/model_training/full")
+    run_inference("examples/qwen_image/model_training/validate_full")
+    run_train_single_GPU("examples/qwen_image/model_training/lora")
+    run_inference("examples/qwen_image/model_training/validate_lora")
+    
+
+def test_wan():
+    run_train_single_GPU("examples/wanvideo/model_inference")
+    move_files("video_", "data/output/model_inference")
+    run_train_single_GPU("examples/wanvideo/model_inference_low_vram")
+    move_files("video_", "data/output/model_inference_low_vram")
+    run_train_multi_GPU("examples/wanvideo/model_training/full")
+    run_train_single_GPU("examples/wanvideo/model_training/validate_full")
+    move_files("video_", "data/output/validate_full")
+    run_train_single_GPU("examples/wanvideo/model_training/lora")
+    run_train_single_GPU("examples/wanvideo/model_training/validate_lora")
+    move_files("video_", "data/output/validate_lora")
+
+
+def test_flux():
+    run_inference("examples/flux/model_inference")
+    run_inference("examples/flux/model_inference_low_vram")
+    run_train_multi_GPU("examples/flux/model_training/full")
+    run_inference("examples/flux/model_training/validate_full")
+    run_train_single_GPU("examples/flux/model_training/lora")
+    run_inference("examples/flux/model_training/validate_lora")
+
+
+def test_z_image():
+    run_inference("examples/z_image/model_inference")
+    run_inference("examples/z_image/model_inference_low_vram")
+    run_train_multi_GPU("examples/z_image/model_training/full")
+    run_inference("examples/z_image/model_training/validate_full")
+    run_train_single_GPU("examples/z_image/model_training/lora")
+    run_inference("examples/z_image/model_training/validate_lora")
+
+
+if __name__ == "__main__":
+    test_z_image()

examples/dev_tools/webui.py

@@ -0,0 +1,408 @@
+import importlib, inspect, pkgutil, traceback, torch, os, re, typing, io
+from typing import Union, List, Optional, Tuple, Iterable, Dict, Literal
+from contextlib import contextmanager
+from diffsynth.utils.data import VideoData
+import streamlit as st
+from diffsynth import ModelConfig
+from diffsynth.diffusion.base_pipeline import ControlNetInput
+from PIL import Image
+from tqdm import tqdm
+st.set_page_config(layout="wide")
+
+class StreamlitTqdmWrapper:
+    """Wrapper class that combines tqdm and streamlit progress bar"""
+    def __init__(self, iterable, st_progress_bar=None):
+        self.iterable = iterable
+        self.st_progress_bar = st_progress_bar
+        self.tqdm_bar = tqdm(iterable)
+        self.total = len(iterable) if hasattr(iterable, '__len__') else None
+        self.current = 0
+
+    def __iter__(self):
+        for item in self.tqdm_bar:
+            if self.st_progress_bar is not None and self.total is not None:
+                self.current += 1
+                self.st_progress_bar.progress(self.current / self.total)
+            yield item
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        if hasattr(self.tqdm_bar, '__exit__'):
+            self.tqdm_bar.__exit__(*args)
+
+@contextmanager
+def catch_error(error_value):
+    try:
+        yield
+    except Exception as e:
+        error_message = traceback.format_exc()
+        print(f"Error {error_value}:\n{error_message}")
+
+def parse_vram_config_from_an_example(path):
+    vram_config = {
+        "offload_dtype": None,
+        "offload_device": None,
+        "onload_dtype": None,
+        "onload_device": None,
+        "preparing_dtype": None,
+        "preparing_device": None,
+        "computation_dtype": None,
+        "computation_device": None,
+    }
+    with open(path, "r") as f:
+        for code in f.readlines():
+            code = code.strip()
+            for param in vram_config:
+                if vram_config[param] is None and f'"{param}":' in code:
+                    value = code.split(" ")[-1].replace(",", "").replace('"', "").replace("torch.", "")
+                    vram_config[param] = value
+    return vram_config
+
+def parse_model_configs_from_an_example(path):
+    model_configs = []
+    vram_config = parse_vram_config_from_an_example(path)
+    with open(path, "r") as f:
+        for code in f.readlines():
+            code = code.strip()
+            if not code.startswith("ModelConfig"):
+                continue
+            pairs = re.findall(r'(\w+)\s*=\s*["\']([^"\']+)["\']', code)
+            config_dict = {k: v for k, v in pairs}
+            vram_config_ = vram_config if "**vram_config" in code else {}
+            model_configs.append(ModelConfig(model_id=config_dict["model_id"], origin_file_pattern=config_dict["origin_file_pattern"], **vram_config_))
+    return model_configs
+
+def list_examples(path, keyword=None):
+    examples = []
+    if os.path.isdir(path):
+        for file_name in os.listdir(path):
+            examples.extend(list_examples(os.path.join(path, file_name), keyword=keyword))
+    elif path.endswith(".py"):
+        with open(path, "r") as f:
+            code = f.read()
+        if keyword is None or keyword in code:
+            examples.extend([path])
+    return examples
+
+def parse_available_pipelines():
+    from diffsynth.diffusion.base_pipeline import BasePipeline
+    import diffsynth.pipelines as _pipelines_pkg
+    available_pipelines = {}
+    for _, name, _ in pkgutil.iter_modules(_pipelines_pkg.__path__):
+        with catch_error(f"Failed: import diffsynth.pipelines.{name}"):
+            mod = importlib.import_module(f"diffsynth.pipelines.{name}")
+            classes = {
+                cls_name: cls for cls_name, cls in inspect.getmembers(mod, inspect.isclass)
+                if issubclass(cls, BasePipeline) and cls is not BasePipeline and cls.__module__ == mod.__name__
+            }
+            available_pipelines.update(classes)
+    return available_pipelines
+
+def parse_available_examples(path, available_pipelines):
+    available_examples = {}
+    for pipeline_name in available_pipelines:
+        examples = ["None"] + list_examples(path, keyword=f"{pipeline_name}.from_pretrained")
+        available_examples[pipeline_name] = examples
+    return available_examples
+
+def draw_selectbox(label, options, option_map, value=None, disabled=False):
+    default_index = 0 if value is None else tuple(options).index([option for option in option_map if option_map[option]==value][0])
+    option = st.selectbox(label=label, options=tuple(options), index=default_index, disabled=disabled)
+    return option_map.get(option)
+
+def parse_params(fn):
+    params = []
+    for name, param in inspect.signature(fn).parameters.items():
+        annotation = param.annotation if param.annotation is not inspect.Parameter.empty else None
+        default = param.default if param.default is not inspect.Parameter.empty else None
+        params.append({"name": name, "dtype": annotation, "value": default})
+    return params
+
+def draw_vram_device(label, value=None, key_suffix="", disabled=False):
+    option_map = {"None": None, "disk": "disk", "cuda": "cuda", "cpu": "cpu"}
+    options = option_map.keys()
+    default_index = 0 if value is None else tuple(options).index(value)
+    option = st.selectbox(label=label, options=tuple(options), index=default_index, key=label + key_suffix, disabled=disabled)
+    return option_map.get(option)
+
+def draw_vram_dtype(label, value=None, key_suffix="", disabled=False):
+    option_map = {"None": None, "disk": "disk", "bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16, "float8_e4m3fn": torch.float8_e4m3fn, "float8_e5m2": torch.float8_e5m2}
+    options = option_map.keys()
+    default_index = 0 if value is None else tuple(options).index(value)
+    option = st.selectbox(label=label, options=tuple(options), index=default_index, key=label + key_suffix, disabled=disabled)
+    return option_map.get(option)
+
+def draw_model_config(model_config=None, key_suffix="", disabled=False, enable_vram_config=False):
+    with st.container(border=True):
+        if model_config is None:
+            model_config = ModelConfig()
+        path = st.text_input(label="path", key="path" + key_suffix, value=model_config.path, disabled=disabled)
+        col1, col2 = st.columns(2)
+        with col1:
+            model_id = st.text_input(label="model_id", key="model_id" + key_suffix, value=model_config.model_id, disabled=disabled)
+        with col2:
+            origin_file_pattern = st.text_input(label="origin_file_pattern", key="origin_file_pattern" + key_suffix, value=model_config.origin_file_pattern, disabled=disabled)
+        if enable_vram_config:
+            with st.container(border=True):
+                col1, col2 = st.columns(2)
+                with col1:
+                    offload_device = draw_vram_device(label="offload_device", value=model_config.offload_device, key_suffix=key_suffix, disabled=disabled)
+                    onload_device = draw_vram_device(label="onload_device", value=model_config.onload_device, key_suffix=key_suffix, disabled=disabled)
+                    preparing_device = draw_vram_device(label="preparing_device", value=model_config.preparing_device, key_suffix=key_suffix, disabled=disabled)
+                    computation_device = draw_vram_device(label="computation_device", value=model_config.computation_device, key_suffix=key_suffix, disabled=disabled)
+                with col2:
+                    offload_dtype = draw_vram_dtype(label="offload_dtype", value=model_config.offload_dtype, key_suffix=key_suffix, disabled=disabled)
+                    onload_dtype = draw_vram_dtype(label="onload_dtype", value=model_config.onload_dtype, key_suffix=key_suffix, disabled=disabled)
+                    preparing_dtype = draw_vram_dtype(label="preparing_dtype", value=model_config.preparing_dtype, key_suffix=key_suffix, disabled=disabled)
+                    computation_dtype = draw_vram_dtype(label="computation_dtype", value=model_config.computation_dtype, key_suffix=key_suffix, disabled=disabled)
+                vram_config = {
+                    "offload_device": offload_device,
+                    "onload_device": onload_device,
+                    "preparing_device": preparing_device,
+                    "computation_device": computation_device,
+                    "offload_dtype": offload_dtype,
+                    "onload_dtype": onload_dtype,
+                    "preparing_dtype": preparing_dtype,
+                    "computation_dtype": computation_dtype,
+                }
+        else:
+            vram_config = {}
+        model_config = ModelConfig(
+            path=None if path == "" else path,
+            model_id=model_id,
+            origin_file_pattern=origin_file_pattern,
+            **vram_config,
+        )
+    return model_config
+
+def draw_multi_model_config(name="", value=None, disabled=False, enable_vram_config=False):
+    model_configs = []
+    with st.container(border=True):
+        st.markdown(name)
+        num = st.number_input(f"num_{name}", min_value=0, max_value=20, value=0 if value is None else len(value), disabled=disabled)
+        for i in range(num):
+            model_config = draw_model_config(key_suffix=f"_{name}_{i}", model_config=None if value is None else value[i], disabled=disabled, enable_vram_config=enable_vram_config)
+            model_configs.append(model_config)
+    return model_configs
+
+def draw_single_model_config(name="", value=None, disabled=False):
+    with st.container(border=True):
+        st.markdown(name)
+        model_config = draw_model_config(value, key_suffix=f"_{name}", disabled=disabled)
+    return model_config
+
+def draw_multi_images(name="", value=None, disabled=False):
+    images = []
+    with st.container(border=True):
+        st.markdown(name)
+        num = st.number_input(f"num_{name}", min_value=0, max_value=20, value=0 if value is None else len(value), disabled=disabled)
+        for i in range(num):
+            image = st.file_uploader(name, type=["png", "jpg", "jpeg", "webp"], key=f"{name}_{i}", disabled=disabled)
+            if image is not None: images.append(Image.open(image))
+    return images
+
+def draw_multi_elements(st_element, name="", value=None, disabled=False, kwargs=None):
+    if kwargs is None:
+        kwargs = {}
+    elements = []
+    with st.container(border=True):
+        st.markdown(name)
+        num = st.number_input(f"num_{name}", min_value=0, max_value=20, value=0 if value is None else len(value), disabled=disabled)
+        for i in range(num):
+            element = st_element(name, key=f"{name}_{i}", disabled=disabled, value=None if value is None else value[i], **kwargs)
+            elements.append(element)
+    return elements
+
+def draw_lora_configs(name="", value=None, disabled=False):
+    elements = []
+    with st.container(border=True):
+        st.markdown(name)
+        num = st.number_input(f"num_{name}", min_value=0, max_value=20, value=0 if value is None else len(value), disabled=disabled)
+        for i in range(num):
+            with st.container(border=True):
+                lora_base_model = st.text_input(label="LoRA base model", key="LoRA base model" + f"LoRA_{i}")
+                lora_scale = st.slider(label="LoRA scale", min_value=-8.0, max_value=8.0, value=1.0, step=0.1, key="LoRA scale" + f"LoRA_{i}")
+                lora_config = draw_model_config(key_suffix=f"LoRA_{i}", disabled=disabled)
+                element = {"base_model": lora_base_model, "alpha": lora_scale, "lora_config": lora_config}
+                elements.append(element)
+    return elements
+
+def draw_controlnet_input(name="", value=None, disabled=False):
+    with st.container(border=True):
+        st.markdown(name)
+        controlnet_id = st.number_input("controlnet_id", value=0, min_value=0, max_value=20, step=1, key=f"{name}_controlnet_id")
+        scale = st.number_input("scale", value=1.0, min_value=0.0, max_value=10.0, key=f"{name}_scale")
+        image = st.file_uploader("image", type=["png", "jpg", "jpeg", "webp"], disabled=disabled, key=f"{name}_image")
+        if image is not None: image = Image.open(image)
+        inpaint_image = st.file_uploader("inpaint_image", type=["png", "jpg", "jpeg", "webp"], disabled=disabled, key=f"{name}_inpaint_image")
+        if inpaint_image is not None: inpaint_image = Image.open(inpaint_image)
+        inpaint_mask = st.file_uploader("inpaint_mask", type=["png", "jpg", "jpeg", "webp"], disabled=disabled, key=f"{name}_inpaint_mask")
+        if inpaint_mask is not None: inpaint_mask = Image.open(inpaint_mask)
+    return ControlNetInput(controlnet_id=controlnet_id, scale=scale, image=image, inpaint_image=inpaint_image, inpaint_mask=inpaint_mask)
+
+def draw_controlnet_inputs(name, value=None, disabled=False):
+    controlnet_inputs = []
+    with st.container(border=True):
+        st.markdown(name)
+        num = st.number_input(f"num_{name}", min_value=0, max_value=20, value=0 if value is None else len(value), disabled=disabled)
+        for i in range(num):
+            controlnet_input = draw_controlnet_input(name=f"{name}_{i}", value=None, disabled=disabled)
+            controlnet_inputs.append(controlnet_input)
+    return controlnet_inputs
+
+def draw_ui_element(name, dtype, value):
+    unsupported_dtype = [
+        Dict[str, torch.Tensor],
+        torch.Tensor,
+    ]
+    if dtype in unsupported_dtype:
+        return
+    if value is None:
+        with st.container(border=True):
+            enable = st.checkbox(f"Enable {name}", value=False)
+            ui = draw_ui_element_safely(name, dtype, value=value, disabled=not enable)
+        if enable:
+            return ui
+        else:
+            return None
+    else:
+        return draw_ui_element_safely(name, dtype, value)
+
+def draw_video(name, value=None, disabled=False):
+    ui = st.file_uploader(name, type=["mp4"], disabled=disabled)
+    if ui is not None:
+        ui = VideoData(ui)
+        ui = [ui[i] for i in range(len(ui))]
+    return ui
+
+def draw_ui_element_safely(name, dtype, value, disabled=False):
+    if dtype == torch.dtype:
+        option_map = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}
+        ui = draw_selectbox(name, option_map.keys(), option_map, value=value, disabled=disabled)
+    elif dtype == Union[str, torch.device]:
+        option_map = {"cuda": "cuda", "cpu": "cpu"}
+        ui = draw_selectbox(name, option_map.keys(), option_map, value=value, disabled=disabled)
+    elif dtype == bool:
+        ui = st.checkbox(name, value=value, disabled=disabled)
+    elif dtype == ModelConfig:
+        ui = draw_single_model_config(name, value=value, disabled=disabled)
+    elif dtype in [list[ModelConfig], List[ModelConfig], Union[list[ModelConfig], ModelConfig, str]]:
+        if name == "model_configs":
+            model_configs = st.session_state.get("model_configs_from_example")
+            ui = draw_multi_model_config(name, model_configs, disabled=disabled, enable_vram_config=True)
+        else:
+            ui = draw_multi_model_config(name, disabled=disabled)
+    elif dtype == str:
+        if "prompt" in name:
+            ui = st.text_area(name, value=value, height=3, disabled=disabled)
+        else:
+            ui = st.text_input(name, value=value, disabled=disabled)
+    elif dtype == float:
+        ui = st.number_input(name, value=value, disabled=disabled)
+    elif dtype == int:
+        ui = st.number_input(name, value=value, step=1, disabled=disabled)
+    elif dtype == Image.Image:
+        ui = st.file_uploader(name, type=["png", "jpg", "jpeg", "webp"], disabled=disabled)
+        if ui is not None: ui = Image.open(ui)
+    elif dtype in [List[Image.Image], list[Image.Image], Union[list[Image.Image], Image.Image], Union[List[Image.Image], Image.Image]]:
+        if "video" in name:
+            ui = draw_video(name, value=value, disabled=disabled)
+        else:
+            ui = draw_multi_images(name, value=value, disabled=disabled)
+    elif dtype in [List[ControlNetInput], list[ControlNetInput]]:
+        ui = draw_controlnet_inputs(name, value=value, disabled=disabled)
+    elif dtype in [List[str], list[str]]:
+        ui = draw_multi_elements(st.text_input, name, value=value, disabled=disabled)
+    elif dtype in [List[float], list[float], Union[list[float], float], Union[List[float], float]]:
+        ui = draw_multi_elements(st.number_input, name, value=value, disabled=disabled)
+    elif dtype in [List[int], list[int]]:
+        ui = draw_multi_elements(st.number_input, name, value=value, disabled=disabled, kwargs={"step": 1})
+    elif dtype in [List[List[Image.Image]], list[list[Image.Image]]]:
+        ui = draw_multi_elements(draw_video, name, value=value, disabled=disabled)
+    elif dtype in [tuple[int, int], Tuple[int, int]]:
+        with st.container(border=True):
+            st.markdown(name)
+            ui = (st.text_input(f"{name}_0", value=value[0], disabled=disabled), st.text_input(f"{name}_1", value=value[1], disabled=disabled))
+    elif isinstance(dtype, typing._LiteralGenericAlias):
+        with st.container(border=True):
+            st.markdown(f"{name} ({dtype})")
+            ui = st.text_input(name, value=value, disabled=disabled, label_visibility="hidden")
+    elif dtype is None:
+        if name == "progress_bar_cmd":
+            ui = value
+    else:
+        st.markdown(f"(`{name}` is not not configurable in WebUI). dtype: `{dtype}`.")
+        ui = value
+    return ui
+
+def flush_example():
+    for key in list(st.session_state.keys()):
+        if key not in ["available_pipelines", "available_examples"]:
+            del st.session_state[key]
+
+def launch_webui():
+    input_col, output_col = st.columns(2)
+    with input_col:
+        if "available_pipelines" not in st.session_state:
+            st.session_state["available_pipelines"] = parse_available_pipelines()
+        if "available_examples" not in st.session_state:
+            st.session_state["available_examples"] = parse_available_examples("./examples", st.session_state["available_pipelines"])
+
+        with st.expander("Pipeline", expanded=True):
+            pipeline_class = draw_selectbox("Pipeline Class", st.session_state["available_pipelines"].keys(), st.session_state["available_pipelines"], value=st.session_state["available_pipelines"]["ZImagePipeline"])
+            example = st.selectbox("Parse model configs from an example (optional)", st.session_state["available_examples"][pipeline_class.__name__], on_change=flush_example)
+
+        if st.button("Step 1: Parse Pipeline", type="primary"):
+            st.session_state["pipeline_class"] = pipeline_class
+            if example != "None":
+                st.session_state["model_configs_from_example"] = parse_model_configs_from_an_example(example)
+
+        if "pipeline_class" not in st.session_state:
+            return
+        with st.expander("Model", expanded=True):
+            input_params = {}
+            params = parse_params(pipeline_class.from_pretrained)
+            for param in params:
+                input_params[param["name"]] = draw_ui_element(**param)
+            lora_configs = draw_lora_configs(name="LoRA")
+        if st.button("Step 2: Load Models", type="primary"):
+            with st.spinner("Loading models", show_time=True):
+                if "pipe" in st.session_state:
+                    del st.session_state["pipe"]
+                    torch.cuda.empty_cache()
+                pipe = pipeline_class.from_pretrained(**input_params)
+                for lora_config in lora_configs:
+                    pipe.load_lora(pipe.get_module(pipe, lora_config["base_model"]), lora_config=lora_config["lora_config"], alpha=lora_config["alpha"])
+                st.session_state["pipe"] = pipe
+
+        if "pipe" not in st.session_state:
+            return
+        with st.expander("Input", expanded=True):
+            pipe = st.session_state["pipe"]
+            input_params = {}
+            params = parse_params(pipeline_class.__call__)
+            for param in params:
+                if param["name"] in ["self"]:
+                    continue
+                input_params[param["name"]] = draw_ui_element(**param)
+    
+    with output_col:
+        if st.button("Step 3: Generate", type="primary"):
+            if "progress_bar_cmd" in input_params:
+                input_params["progress_bar_cmd"] = lambda iterable: StreamlitTqdmWrapper(iterable, st.progress(0))
+            result = pipe(**input_params)
+            st.session_state["result"] = result
+        
+        if "result" in st.session_state:
+            result = st.session_state["result"]
+            if isinstance(result, Image.Image):
+                st.image(result)
+                buf = io.BytesIO()
+                result.save(buf, format='PNG')
+                st.download_button(label="Download", data=buf.getvalue(), file_name="image.png", mime="image/png", type="primary")
+            else:
+                print(f"unsupported result format: {result}")
+
+launch_webui()

examples/dev_tools/webui_train.py

@@ -0,0 +1,445 @@
+import os, importlib.util, argparse, pkgutil, inspect
+from dataclasses import dataclass
+import streamlit as st
+st.set_page_config(layout="wide")
+
+available_data_file_keys = ["animate_face_video", "animate_pose_video", "audio", "blockwise_controlnet_image", "blockwise_controlnet_inpaint_mask", "context_image", "control_video", "controlnet_image", "controlnet_inpaint_mask", "edit_image", "eligen_entity_masks", "image", "in_context_videos", "infinityou_id_image", "input_audio", "ipadapter_images", "kontext_images", "layer_input_image", "nexus_gen_reference_image", "reference_image", "s2v_pose_video", "step1x_reference_image", "vace_reference_image", "vace_video", "vap_video", "video", "wantodance_keyframes", "wantodance_music_path", "wantodance_reference_image"]
+available_extra_inputs = ["animate_face_video", "animate_pose_video", "blockwise_controlnet_image", "blockwise_controlnet_inpaint_mask", "camera_control_direction", "camera_control_speed", "cfg_scale", "context_image", "control_video", "controlnet_image", "controlnet_inpaint_mask", "controlnet_processor_id", "edit_image", "eligen_entity_masks", "eligen_entity_prompts", "end_image", "frame_rate", "framewise_decoding", "in_context_downsample_factor", "in_context_videos", "infinityou_guidance", "infinityou_id_image", "input_audio", "input_image", "ipadapter_images", "kontext_images", "layer_input_image", "layer_num", "lora_encoder_inputs", "motion_bucket_id", "nexus_gen_reference_image", "num_inference_steps", "rand_device", "reference_image", "s2v_pose_video", "seed", "step1x_reference_image", "template_inputs", "vace_reference_image", "vace_video", "value_controller_inputs", "vap_video", "wantodance_fps", "wantodance_keyframes", "wantodance_keyframes_mask", "wantodance_music_path", "wantodance_reference_image"]
+available_model_components = ["animate_adapter", "audio_dit", "audio_encoder", "audio_vae", "audio_vae_decoder", "audio_vae_encoder", "audio_vocoder", "blockwise_controlnet", "conditioner", "controlnet", "dinov3_image_encoder", "dit", "dit2", "dual_tower_bridge", "image2lora_coarse", "image2lora_fine", "image2lora_style", "image_encoder", "image_proj_model", "infinityou_processor", "ipadapter", "ipadapter_image_encoder", "lora_encoder", "lora_patcher", "motion_controller", "nexus_gen", "nexus_gen_editing_adapter", "nexus_gen_generation_adapter", "qwenvl", "siglip2_image_encoder", "step1x_connector", "text_encoder", "text_encoder_1", "text_encoder_2", "text_encoder_post_modules", "text_encoder_qwen3", "tokenizer_model", "tokenizer_t5xxl", "unet", "upsampler", "vace", "vace2", "vae", "vae_decoder", "vae_encoder", "value_controller", "vap", "video_dit", "video_dit2", "video_vae", "video_vae_decoder", "video_vae_encoder"]
+
+@dataclass
+class Parameter:
+    name: str = None
+    dtype: type = None
+    value: any = None
+    required: bool = False
+    choices: list = None
+    help: str = None
+
+def parse_available_pipelines():
+    from diffsynth.diffusion.base_pipeline import BasePipeline
+    import diffsynth.pipelines as _pipelines_pkg
+    available_pipelines = {}
+    for _, name, _ in pkgutil.iter_modules(_pipelines_pkg.__path__):
+        mod = importlib.import_module(f"diffsynth.pipelines.{name}")
+        classes = {
+            cls_name: cls for cls_name, cls in inspect.getmembers(mod, inspect.isclass)
+            if issubclass(cls, BasePipeline) and cls is not BasePipeline and cls.__module__ == mod.__name__
+        }
+        available_pipelines.update(classes)
+    return available_pipelines
+
+def search_for_options(name):
+    files = search_for_files("examples", ".sh")
+    params = set()
+    for file in files:
+        with open(file, "r", encoding="utf-8") as f:
+            for line in f.readlines():
+                if f"--{name}" in line and not line.startswith("#"):
+                    line = line.strip()
+                    line = line.replace(" \\", "").replace(f"--{name}", "").replace('"', "").replace(" ", "")
+                    for param in line.split(","):
+                        params.add(param)
+    for param in sorted(list(params)):
+        print(f'"{param}", ', end="")
+    print()
+
+def search_for_available_pipeline_options():
+    from diffsynth.diffusion.base_pipeline import BasePipeline
+    pipeline_classes = parse_available_pipelines()
+    base_attrs = set(vars(BasePipeline()))
+    black_list = ["tokenizer", "processor", "tokenizer_1", "tokenizer_2", "audio_processor"]
+    options = []
+    for pipeline_class in pipeline_classes:
+        pipe = pipeline_classes[pipeline_class]()
+        members = [attr for attr in vars(pipe) if not attr.startswith("__") and attr not in base_attrs and getattr(pipe, attr) is None]
+        members = [attr for attr in members if attr not in black_list]
+        options.extend(members)
+    options = sorted(list(set(options)))
+    for option in sorted(list(options)):
+        print(f'"{option}", ', end="")
+    print()
+
+def parse_available_training_scripts(path):
+    training_scripts = {}
+    for folder in os.listdir(path):
+        if os.path.isfile(f"{path}/{folder}/model_training/train.py"):
+            training_scripts[folder] = f"{path}/{folder}/model_training/train.py"
+    return training_scripts
+
+def search_for_files(path, suffix):
+    if os.path.isfile(path):
+        if path.endswith(suffix): return [path]
+        return []
+    else:
+        files = []
+        for sub_path in os.listdir(path):
+            files.extend(search_for_files(os.path.join(path, sub_path), suffix))
+        return files
+
+def parse_available_examples(path):
+    path = os.path.dirname(path)
+    examples = search_for_files(path, ".sh")
+    return examples
+
+def parse_example(example_path):
+    value_dict = {}
+    with open(example_path, "r", encoding="utf-8") as f:
+        for line in f.readlines():
+            line = line.strip()
+            if line.startswith("#"):
+                continue
+            if not line.startswith("--"):
+                continue
+            line = line.replace("\\", "").strip()
+            if " " in line:
+                name, value = line[2:line.index(" ")], line[line.index(" ") + 1:]
+                if value.startswith('"') and value.endswith('"'):
+                    value = value[1:-1]
+                if value.startswith("'") and value.endswith("'"):
+                    value = value[1:-1]
+            else:
+                name, value = line[2:], True
+            value_dict[name] = value
+    return value_dict
+
+def parse_parser(path):
+    spec = importlib.util.spec_from_file_location("train", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    for name in dir(module):
+        if name.endswith("parser") and callable(getattr(module, name)):
+            return getattr(module, name)
+    return None
+
+def parse_parser_action(action, value=None):
+    if isinstance(action, argparse._StoreTrueAction) or isinstance(action, argparse._StoreFalseAction):
+        dtype = bool
+    else:
+        dtype = action.type
+    param = Parameter(
+        name=action.dest,
+        dtype=dtype,
+        value=action.default if value is None else dtype(value),
+        required=action.required,
+        choices=action.choices,
+        help=action.help,
+    )
+    return param
+
+def parse_parser_actions(parser, example_path=None):
+    value_dict = {} if example_path is None or example_path == "None" else parse_example(example_path)
+    params = []
+    for action in parser._actions:
+        param = parse_parser_action(action, value=value_dict.get(action.dest))
+        if param.name == "help":
+            continue
+        params.append(param)
+    return params
+
+def draw_model_id_with_origin_paths(param, disabled=False):
+    with st.container(border=True):
+        st.markdown(param.name, help=param.help)
+        model_id_with_origin_paths = [] if param.value is None else param.value.split(",")
+        num = st.number_input(f"Number of models", min_value=0, max_value=20, value=len(model_id_with_origin_paths), disabled=disabled)
+        result = []
+        for i in range(num):
+            col1, col2 = st.columns(2)
+            value = model_id_with_origin_paths[i].split(":") if i < len(model_id_with_origin_paths) else (None, None)
+            with col1:
+                model_id = st.text_input("model_id", value=value[0], key=f"model_id_{i}", disabled=disabled)
+            with col2:
+                origin_file_pattern = st.text_input("origin_file_pattern", value=value[1], key=f"origin_file_pattern_{i}", disabled=disabled)
+            result.append(f"{model_id}:{origin_file_pattern}")
+        result = ",".join(result)
+    return result
+
+def draw_parameter(param, check_enable=True, disabled=False):
+    if check_enable and param.value is None:
+        with st.container(border=True):
+            enable_button = st.checkbox(f"Enable {param.name}", value=False, disabled=disabled)
+            ui = draw_parameter(param, check_enable=False, disabled=disabled or not enable_button)
+            if enable_button:
+                return ui
+            else:
+                return None
+    if param.name == "data_file_keys":
+        ui = st.multiselect(param.name, options=available_data_file_keys, accept_new_options=True, default=param.value.split(","), disabled=disabled, help=param.help)
+        ui = ",".join(ui)
+    elif param.name == "model_paths":
+        ui = st.text_area(param.name, value=param.value, height=3, disabled=disabled, help=param.help)
+    elif param.name == "model_id_with_origin_paths":
+        ui = draw_model_id_with_origin_paths(param, disabled=disabled)
+    elif param.name == "extra_inputs":
+        value = None if param.value is None else param.value.split(",")
+        ui = st.multiselect(param.name, options=available_extra_inputs, accept_new_options=True, default=value, disabled=disabled, help=param.help)
+        ui = ",".join(ui)
+    elif param.name in ["fp8_models", "offload_models", "trainable_models", "lora_base_model", "preset_lora_model"]:
+        value = None if param.value is None else param.value.split(",")
+        ui = st.multiselect(param.name, options=available_model_components, accept_new_options=True, default=value, disabled=disabled, help=param.help)
+        ui = ",".join(ui)
+    elif param.name == "learning_rate":
+        ui = st.number_input(param.name, value=param.value, format="%0.7f", step=1e-4, disabled=disabled, help=param.help)
+    elif param.dtype == str:
+        ui = st.text_input(param.name, value=param.value, disabled=disabled, help=param.help)
+    elif param.dtype == int:
+        ui = st.number_input(param.name, value=param.value, step=1, disabled=disabled, help=param.help)
+    elif param.dtype == float:
+        ui = st.number_input(param.name, value=param.value, disabled=disabled, help=param.help)
+    elif param.dtype == bool:
+        ui = st.checkbox(param.name, value=param.value, disabled=disabled, help=param.help)
+    else:
+        st.markdown(f"(`{param.name}` is not not configurable in WebUI). dtype: `{param.dtype}`.")
+        ui = None
+    return ui
+
+def draw_dataset_configs(dataset_base_path, dataset_metadata_path, dataset_repeat, dataset_num_workers, data_file_keys):
+    dataset_base_path = draw_parameter(dataset_base_path)
+    dataset_metadata_path = draw_parameter(dataset_metadata_path)
+    col_1, col_2 = st.columns(2)
+    with col_1:
+        dataset_repeat = draw_parameter(dataset_repeat)
+    with col_2:
+        dataset_num_workers = draw_parameter(dataset_num_workers)
+    data_file_keys = draw_parameter(data_file_keys)
+    params = {
+        "dataset_base_path": dataset_base_path,
+        "dataset_metadata_path": dataset_metadata_path,
+        "dataset_repeat": dataset_repeat,
+    }
+    if dataset_num_workers > 0: params["dataset_num_workers"] = dataset_num_workers
+    params["data_file_keys"] = data_file_keys
+    return params
+
+def draw_image_size(height, width, max_pixels):
+    mode = st.selectbox("Image scaling and cropping", options=["Scale if pixel count exceeds threshold", "Resize to a fixed size and crop"], index=int(max_pixels.value is None))
+    use_max_pixel = mode == "Scale if pixel count exceeds threshold"
+    use_height_width = mode == "Resize to a fixed size and crop"
+    col1, col2 = st.columns(2)
+    with col1:
+        with st.container(border=True):
+            max_pixels = draw_parameter(max_pixels, disabled=not use_max_pixel)
+    with col2:
+        with st.container(border=True):
+            height = draw_parameter(height, disabled=not use_height_width, check_enable=False)
+            width = draw_parameter(width, disabled=not use_height_width, check_enable=False)
+    if use_max_pixel:
+        return {"max_pixels": max_pixels}
+    else:
+        return {"height": height, "width": width}
+
+def draw_model_configs(model_paths, model_id_with_origin_paths, extra_inputs, fp8_models, offload_models):
+    model_id_with_origin_paths = draw_parameter(model_id_with_origin_paths)
+    model_paths = draw_parameter(model_paths)
+    extra_inputs = draw_parameter(extra_inputs)
+    fp8_models = draw_parameter(fp8_models)
+    offload_models = draw_parameter(offload_models)
+    params = {}
+    if model_paths is not None: params["model_paths"] = model_paths
+    if model_id_with_origin_paths is not None: params["model_id_with_origin_paths"] = model_id_with_origin_paths
+    if extra_inputs is not None: params["extra_inputs"] = extra_inputs
+    if fp8_models is not None: params["fp8_models"] = fp8_models
+    if offload_models is not None: params["offload_models"] = offload_models
+    return params
+
+def draw_video_size(height, width, max_pixels, num_frames):
+    mode = st.selectbox("Video scaling and cropping", options=["Scale if pixel count exceeds threshold", "Resize to a fixed size and crop"], index=int(max_pixels.value is not None))
+    use_max_pixel = mode == "Scale if pixel count exceeds threshold"
+    use_height_width = mode == "Resize to a fixed size and crop"
+    col1, col2 = st.columns(2)
+    with col1:
+        with st.container(border=True):
+            max_pixels = draw_parameter(max_pixels, disabled=not use_max_pixel)
+    with col2:
+        with st.container(border=True):
+            height = draw_parameter(height, disabled=not use_height_width)
+            width = draw_parameter(width, disabled=not use_height_width)
+    num_frames = draw_parameter(num_frames)
+    if use_max_pixel:
+        return {"max_pixels": max_pixels, "num_frames": num_frames}
+    else:
+        return {"height": height, "width": width, "num_frames": num_frames}
+
+def draw_training_configs(learning_rate, num_epochs, trainable_models, find_unused_parameters, weight_decay, task):
+    learning_rate = draw_parameter(learning_rate)
+    num_epochs = draw_parameter(num_epochs)
+    trainable_models = draw_parameter(trainable_models)
+    weight_decay = draw_parameter(weight_decay)
+    task = draw_parameter(task)
+    find_unused_parameters = draw_parameter(find_unused_parameters)
+    params = {
+        "task": task,
+        "find_unused_parameters": find_unused_parameters,
+        "num_epochs": num_epochs,
+        "learning_rate": learning_rate,
+    }
+    if weight_decay != 0.01: params["weight_decay"] = weight_decay
+    if trainable_models is not None: params["trainable_models"] = trainable_models
+    return params
+
+def draw_output_configs(output_path, remove_prefix_in_ckpt, save_steps):
+    output_path = draw_parameter(output_path)
+    remove_prefix_in_ckpt = draw_parameter(remove_prefix_in_ckpt)
+    save_steps = draw_parameter(save_steps)
+    params = {
+        "output_path": output_path,
+        "remove_prefix_in_ckpt": remove_prefix_in_ckpt,
+    }
+    if save_steps is not None: params["save_steps"] = save_steps
+    return params
+
+def draw_lora_configs(lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, preset_lora_path, preset_lora_model):
+    with st.container(border=True):
+        train_lora = st.checkbox("Train LoRA", value=lora_base_model is not None)
+        lora_base_model = draw_parameter(lora_base_model, check_enable=False, disabled=not train_lora)
+        lora_target_modules = draw_parameter(lora_target_modules, check_enable=False, disabled=not train_lora)
+        lora_rank = draw_parameter(lora_rank, check_enable=False, disabled=not train_lora)
+        lora_checkpoint = draw_parameter(lora_checkpoint, check_enable=False, disabled=not train_lora)
+    preset_lora_path = draw_parameter(preset_lora_path)
+    preset_lora_model = draw_parameter(preset_lora_model)
+    params = {}
+    if train_lora:
+        if lora_base_model is not None: params["lora_base_model"] = lora_base_model
+        if lora_target_modules is not None: params["lora_target_modules"] = lora_target_modules
+        if lora_rank is not None: params["lora_rank"] = lora_rank
+        if lora_checkpoint is not None: params["lora_checkpoint"] = lora_checkpoint
+    if preset_lora_path is not None: params["preset_lora_path"] = preset_lora_path
+    if preset_lora_model is not None: params["preset_lora_model"] = preset_lora_model
+    return params
+
+def draw_gradient_configs(use_gradient_checkpointing, use_gradient_checkpointing_offload, gradient_accumulation_steps):
+    use_gradient_checkpointing = draw_parameter(use_gradient_checkpointing)
+    use_gradient_checkpointing_offload = draw_parameter(use_gradient_checkpointing_offload)
+    gradient_accumulation_steps = draw_parameter(gradient_accumulation_steps)
+    params = {
+        "use_gradient_checkpointing": use_gradient_checkpointing,
+        "use_gradient_checkpointing_offload": use_gradient_checkpointing_offload,
+    }
+    if gradient_accumulation_steps != 1:
+        params["gradient_accumulation_steps"] = gradient_accumulation_steps
+    return params
+
+def draw_template_model_configs(template_model_id_or_path, enable_lora_hot_loading):
+    template_model_id_or_path = draw_parameter(template_model_id_or_path)
+    enable_lora_hot_loading = draw_parameter(enable_lora_hot_loading)
+    params = {"enable_lora_hot_loading": enable_lora_hot_loading}
+    if template_model_id_or_path is not None: params["template_model_id_or_path"] = template_model_id_or_path
+    return params
+
+def match_ui_groups(params, ui_groups):
+    param_names = [param.name for param in params]
+    for ui_group in ui_groups:
+        if sum([name in param_names for name in ui_group["params"]]) == len(ui_group["params"]):
+            group_params = {param.name: param for param in params if param.name in ui_group["params"]}
+            other_params = [param for param in params if param.name not in ui_group["params"]]
+            return group_params, other_params, ui_group
+    return {}, params, None
+
+def draw_other_params(params):
+    results = {}
+    for param in params:
+        results[param.name] = draw_parameter(param)
+    return results
+
+def draw_all_params(params, ui_groups):
+    matched_ui_groups = []
+    while True:
+        group_params, other_params, ui_group = match_ui_groups(params, ui_groups)
+        if len(group_params) == 0:
+            break
+        matched_ui_groups.append((ui_group, group_params))
+        params = other_params
+    tabs = st.tabs([ui_group["name"] for ui_group, _ in matched_ui_groups] + ["Others"])
+    inputs = {}
+    for tab, (ui_group, group_params) in zip(tabs, matched_ui_groups):
+        with tab:
+            inputs.update(ui_group["fn"](**group_params))
+    with tabs[-1]:
+        inputs.update(draw_other_params(params))
+    return inputs
+
+def generate_training_script(script_path, inputs):
+    cmd = f"accelerate launch {script_path}"
+    for name, value in inputs.items():
+        if value is not None:
+            if isinstance(value, bool):
+                if value == True:
+                    cmd = f"{cmd} \\\n  --{name}"
+            elif isinstance(value, str):
+                cmd = f"{cmd} \\\n  --{name} \"{value}\""
+            else:
+                cmd = f"{cmd} \\\n  --{name} {value}"
+    return cmd
+
+ui_groups = [
+    {
+        "name": "Dataset",
+        "params": ("dataset_base_path", "dataset_metadata_path", "dataset_repeat", "dataset_num_workers", "data_file_keys"),
+        "fn": draw_dataset_configs,
+    },
+    {
+        "name": "Video Size",
+        "params": ("height", "width", "max_pixels", "num_frames"),
+        "fn": draw_video_size,
+    },
+    {
+        "name": "Image Size",
+        "params": ("height", "width", "max_pixels"),
+        "fn": draw_image_size,
+    },
+    {
+        "name": "Model",
+        "params": ("model_paths", "model_id_with_origin_paths", "extra_inputs", "fp8_models", "offload_models"),
+        "fn": draw_model_configs,
+    },
+    {
+        "name": "Training",
+        "params": ("learning_rate", "num_epochs", "trainable_models", "find_unused_parameters", "weight_decay", "task"),
+        "fn": draw_training_configs,
+    },
+    {
+        "name": "Output",
+        "params": ("output_path", "remove_prefix_in_ckpt", "save_steps"),
+        "fn": draw_output_configs,
+    },
+    {
+        "name": "LoRA",
+        "params": ("lora_base_model", "lora_target_modules", "lora_rank", "lora_checkpoint", "preset_lora_path", "preset_lora_model"),
+        "fn": draw_lora_configs,
+    },
+    {
+        "name": "Gradient",
+        "params": ("use_gradient_checkpointing", "use_gradient_checkpointing_offload", "gradient_accumulation_steps"),
+        "fn": draw_gradient_configs,
+    },
+    {
+        "name": "Templates",
+        "params": ("template_model_id_or_path", "enable_lora_hot_loading"),
+        "fn": draw_template_model_configs,
+    },
+]
+
+def launch_webui():
+    input_col, output_col = st.columns(2)
+    with input_col:
+        if "available_training_scripts" not in st.session_state:
+            st.session_state["available_training_scripts"] = parse_available_training_scripts("examples")
+        with st.container(border=True):
+            script_path = st.selectbox(label="Script path", options=st.session_state["available_training_scripts"].values(), index=0)
+            example_path = st.selectbox(label="Example path (Optional)", options=["None"] + parse_available_examples(script_path), index=0)
+        if st.button("Step 1: Parse Training Script", type="primary"):
+            st.session_state["script_path"] = script_path
+        
+        if "script_path" not in st.session_state:
+            return
+        with st.spinner("Fetching input parameters", show_time=False):
+            parser = parse_parser(script_path)
+        parser = parser()
+        params = parse_parser_actions(parser, example_path)
+        inputs = draw_all_params(params, ui_groups)
+    with output_col:
+        if st.button("Step 2: Generate training script", type="primary"):
+            script = generate_training_script(script_path, inputs)
+            st.code(script, language="shell")
+
+launch_webui()

examples/ernie_image/model_inference/ERNIE-Image-Turbo.py

@@ -0,0 +1,25 @@
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+import torch
+
+pipe = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device='cuda',
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image-Turbo", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="tokenizer/"),
+)
+
+image = pipe(
+    prompt="一只黑白相间的中华田园犬",
+    negative_prompt="",
+    height=1024,
+    width=1024,
+    seed=42,
+    num_inference_steps=8,
+    cfg_scale=1.0,
+    sigma_shift=4.0,
+)
+image.save("output_turbo.jpg")

examples/ernie_image/model_inference/ERNIE-Image.py

@@ -0,0 +1,24 @@
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+import torch
+
+pipe = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device='cuda',
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="tokenizer/"),
+)
+
+image = pipe(
+    prompt="一只黑白相间的中华田园犬",
+    negative_prompt="",
+    height=1024,
+    width=1024,
+    seed=42,
+    num_inference_steps=50,
+    cfg_scale=4.0,
+)
+image.save("output.jpg")

examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py

@@ -0,0 +1,37 @@
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, 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 = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device='cuda',
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image-Turbo", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="tokenizer/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image = pipe(
+    prompt="一只黑白相间的中华田园犬",
+    negative_prompt="",
+    height=1024,
+    width=1024,
+    seed=42,
+    num_inference_steps=8,
+    cfg_scale=1.0,
+    sigma_shift=4.0,
+)
+image.save("output_turbo.jpg")

examples/ernie_image/model_inference_low_vram/ERNIE-Image.py

@@ -0,0 +1,36 @@
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, 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 = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device='cuda',
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="tokenizer/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image = pipe(
+    prompt="一只黑白相间的中华田园犬",
+    negative_prompt="",
+    height=1024,
+    width=1024,
+    seed=42,
+    num_inference_steps=50,
+    cfg_scale=4.0,
+)
+image.save("output.jpg")

examples/ernie_image/model_training/full/ERNIE-Image.sh

@@ -0,0 +1,17 @@
+# Dataset: data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I/
+
+accelerate launch --config_file examples/ernie_image/model_training/full/accelerate_config_zero3.yaml \
+  examples/ernie_image/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I \
+  --dataset_metadata_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "PaddlePaddle/ERNIE-Image:transformer/diffusion_pytorch_model*.safetensors,PaddlePaddle/ERNIE-Image:text_encoder/model.safetensors,PaddlePaddle/ERNIE-Image:vae/diffusion_pytorch_model.safetensors" \
+  --learning_rate 1e-5 \
+  --num_epochs 2 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Ernie-Image-T2I_full" \
+  --trainable_models "dit" \
+  --use_gradient_checkpointing \
+  --dataset_num_workers 8 \
+  --find_unused_parameters

examples/ernie_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/ernie_image/model_training/lora/ERNIE-Image.sh

@@ -0,0 +1,19 @@
+# Dataset: data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I/
+# Download: modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ernie_image/Ernie-Image-T2I/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/ernie_image/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I \
+  --dataset_metadata_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "PaddlePaddle/ERNIE-Image:transformer/diffusion_pytorch_model*.safetensors,PaddlePaddle/ERNIE-Image:text_encoder/model.safetensors,PaddlePaddle/ERNIE-Image:vae/diffusion_pytorch_model.safetensors" \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Ernie-Image-T2I_lora" \
+  --lora_base_model "dit" \
+  --lora_target_modules "to_q,to_k,to_v,to_out.0" \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --dataset_num_workers 8 \
+  --find_unused_parameters

examples/ernie_image/model_training/special/split_training/ERNIE-Image.sh

@@ -0,0 +1,40 @@
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ernie_image/Ernie-Image-T2I/*" --local_dir ./data/diffsynth_example_dataset
+
+# Stage 1: cache deterministic preprocessing outputs.
+accelerate launch examples/ernie_image/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I \
+  --dataset_metadata_path data/diffsynth_example_dataset/ernie_image/Ernie-Image-T2I/metadata.csv \
+  --max_pixels 1048576 \
+  --dataset_repeat 1 \
+  --model_id_with_origin_paths 'PaddlePaddle/ERNIE-Image:transformer/diffusion_pytorch_model*.safetensors,PaddlePaddle/ERNIE-Image:text_encoder/model.safetensors,PaddlePaddle/ERNIE-Image:vae/diffusion_pytorch_model.safetensors' \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/Ernie-Image-T2I_split_cache \
+  --lora_base_model dit \
+  --lora_target_modules to_q,to_k,to_v,to_out.0 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --dataset_num_workers 8 \
+  --find_unused_parameters \
+  --offload_models 'PaddlePaddle/ERNIE-Image:transformer/diffusion_pytorch_model*.safetensors' \
+  --task sft:data_process
+
+# Stage 2: train LoRA from the cached dataset.
+accelerate launch examples/ernie_image/model_training/train.py \
+  --dataset_base_path ./models/train/Ernie-Image-T2I_split_cache \
+  --max_pixels 1048576 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths 'PaddlePaddle/ERNIE-Image:transformer/diffusion_pytorch_model*.safetensors,PaddlePaddle/ERNIE-Image:text_encoder/model.safetensors,PaddlePaddle/ERNIE-Image:vae/diffusion_pytorch_model.safetensors' \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt pipe.dit. \
+  --output_path ./models/train/Ernie-Image-T2I_split \
+  --lora_base_model dit \
+  --lora_target_modules to_q,to_k,to_v,to_out.0 \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --dataset_num_workers 8 \
+  --find_unused_parameters \
+  --offload_models PaddlePaddle/ERNIE-Image:text_encoder/model.safetensors,PaddlePaddle/ERNIE-Image:vae/diffusion_pytorch_model.safetensors \
+  --task sft:train

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

@@ -0,0 +1,25 @@
+import torch
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+from diffsynth.core.loader.file import load_state_dict
+
+pipe = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+lora_state_dict = load_state_dict('./models/train/Ernie-Image-T2I_split/epoch-4.safetensors', torch_dtype=torch.bfloat16, device="cuda")
+pipe.load_lora(pipe.dit, state_dict=lora_state_dict, alpha=1.0)
+
+image = pipe(
+    prompt="a professional photo of a cute dog",
+    seed=0,
+    num_inference_steps=50,
+    cfg_scale=4.0,
+)
+image.save('split_training_ERNIE-Image.jpg')
+print("LoRA validation image saved to image_lora.jpg")

examples/ernie_image/model_training/train.py

@@ -0,0 +1,141 @@
+import torch, os, argparse, accelerate
+from diffsynth.core import UnifiedDataset
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+from diffsynth.diffusion import *
+from diffsynth.core.data.operators import *
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class ErnieImageTrainingModule(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,
+        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)
+        tokenizer_config = ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="tokenizer/") if tokenizer_path is None else ModelConfig(tokenizer_path)
+        self.pipe = ErnieImagePipeline.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.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.task = task
+        self.task_to_loss = {
+            "sft:data_process": lambda pipe, inputs_shared, inputs_posi, inputs_nega: (inputs_shared, inputs_posi, inputs_nega),
+            "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 ernie_image_parser():
+    parser = argparse.ArgumentParser(description="ERNIE-Image 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 = ernie_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=lambda x: x,
+        special_operator_map={
+            "image": ToAbsolutePath(args.dataset_base_path) >> LoadImage() >> ImageCropAndResize(args.height, args.width, args.max_pixels, 16, 16),
+        },
+    )
+    model = ErnieImageTrainingModule(
+        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,
+        resume_from_checkpoint=args.resume_from_checkpoint,
+        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
+        task=args.task,
+        device="cpu" if 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/ernie_image/model_training/validate_full/ERNIE-Image.py

@@ -0,0 +1,25 @@
+import torch
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+from diffsynth.core import load_state_dict
+
+pipe = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+state_dict = load_state_dict("./models/train/Ernie-Image-T2I_full/epoch-1.safetensors")
+pipe.dit.load_state_dict(state_dict)
+
+image = pipe(
+    prompt="a professional photo of a cute dog",
+    seed=0,
+    num_inference_steps=50,
+    cfg_scale=4.0,
+)
+image.save("image_full.jpg")
+print("Full validation image saved to image_full.jpg")

examples/ernie_image/model_training/validate_lora/ERNIE-Image.py

@@ -0,0 +1,25 @@
+import torch
+from diffsynth.pipelines.ernie_image import ErnieImagePipeline, ModelConfig
+from diffsynth.core.loader.file import load_state_dict
+
+pipe = ErnieImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="PaddlePaddle/ERNIE-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+lora_state_dict = load_state_dict("./models/train/Ernie-Image-T2I_lora/epoch-4.safetensors", torch_dtype=torch.bfloat16, device="cuda")
+pipe.load_lora(pipe.dit, state_dict=lora_state_dict, alpha=1.0)
+
+image = pipe(
+    prompt="a professional photo of a cute dog",
+    seed=0,
+    num_inference_steps=50,
+    cfg_scale=4.0,
+)
+image.save("image_lora.jpg")
+print("LoRA validation image saved to image_lora.jpg")

examples/flux/README.md

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

examples/flux/model_inference/FLEX.2-preview.py

@@ -0,0 +1,50 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from diffsynth.utils.controlnet import Annotator
+import numpy as np
+from PIL import Image
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="ostris/Flex.2-preview", origin_file_pattern="Flex.2-preview.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl, long hair, red t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    seed=0
+)
+image.save("image_1.jpg")
+
+mask = np.zeros((1024, 1024, 3), dtype=np.uint8)
+mask[200:400, 400:700] = 255
+mask = Image.fromarray(mask)
+mask.save("image_mask.jpg")
+
+inpaint_image = image
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl with sunglasses, long hair, red t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    flex_inpaint_image=inpaint_image, flex_inpaint_mask=mask,
+    seed=4
+)
+image.save("image_2.jpg")
+
+control_image = Annotator("canny")(image)
+control_image.save("image_control.jpg")
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl with sunglasses, long hair, yellow t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    flex_control_image=control_image,
+    seed=4
+)
+image.save("image_3.jpg")

examples/flux/model_inference/FLUX.1-Fill-dev.py

@@ -0,0 +1,29 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="flux1-fill-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+dataset_snapshot_download(
+    dataset_id="HuanJue/example_dataset",
+    local_dir="./",
+    allow_file_pattern=f"FLUX.1-Fill-dev/*",
+)
+
+flux_fill_image = Image.open("FLUX.1-Fill-dev/cup.png").convert("RGB")
+flux_fill_mask = Image.open("FLUX.1-Fill-dev/cup_mask.png").convert("L")
+prompt = "a white paper cup"
+image = pipe(prompt=prompt, flux_fill_image=flux_fill_image, flux_fill_mask=flux_fill_mask, height=1632, width=1232, seed=0, embedded_guidance=30.0, num_inference_steps=50)
+
+image.save("image_FLUX.1-Fill-dev.jpg")

examples/flux/model_inference/FLUX.1-Kontext-dev.py

@@ -0,0 +1,54 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="flux1-kontext-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+image_1 = pipe(
+    prompt="a beautiful Asian long-haired female college student.",
+    embedded_guidance=2.5,
+    seed=1,
+)
+image_1.save("image_1.jpg")
+
+image_2 = pipe(
+    prompt="transform the style to anime style.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=2,
+)
+image_2.save("image_2.jpg")
+
+image_3 = pipe(
+    prompt="let her smile.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=3,
+)
+image_3.save("image_3.jpg")
+
+image_4 = pipe(
+    prompt="let the girl play basketball.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=4,
+)
+image_4.save("image_4.jpg")
+
+image_5 = pipe(
+    prompt="move the girl to a park, let her sit on a chair.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=5,
+)
+image_5.save("image_5.jpg")

examples/flux/model_inference/FLUX.1-Krea-dev.py

@@ -0,0 +1,27 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Krea-dev", origin_file_pattern="flux1-krea-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+prompt = "An beautiful woman is riding a bicycle in a park, wearing a red dress"
+negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+
+image = pipe(prompt=prompt, seed=0, embedded_guidance=4.5)
+image.save("flux_krea.jpg")
+
+image = pipe(
+    prompt=prompt, negative_prompt=negative_prompt,
+    seed=0, cfg_scale=2, num_inference_steps=50,
+    embedded_guidance=4.5
+)
+image.save("flux_krea_cfg.jpg")

examples/flux/model_inference/FLUX.1-Redux-dev.py

@@ -0,0 +1,28 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_embedder/diffusion_pytorch_model.safetensors"),
+    ],
+)
+dataset_snapshot_download(
+    dataset_id="HuanJue/example_dataset",
+    local_dir="./",
+    allow_file_pattern=f"FLUX.1-Redux-dev/*",
+)
+
+flux_redux_image = Image.open("FLUX.1-Redux-dev/robot.png").convert("RGB")
+
+image = pipe(flux_redux_image=flux_redux_image, embedded_guidance=2.5, num_inference_steps=50)
+
+image.save("image_FLUX.1-Redux-dev.jpg")

examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py

@@ -0,0 +1,19 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors")
+    ],
+)
+
+for i in [0.1, 0.3, 0.5, 0.7, 0.9]:
+    image = pipe(prompt="a cat on the beach", seed=2, value_controller_inputs=[i])
+    image.save(f"value_control_{i}.jpg")

examples/flux/model_inference/FLUX.1-dev-Controlnet-Inpainting-Beta.py

@@ -0,0 +1,37 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+import numpy as np
+from PIL import Image
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", origin_file_pattern="diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+image_1 = pipe(
+    prompt="a cat sitting on a chair",
+    height=1024, width=1024,
+    seed=8, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+mask = np.zeros((1024, 1024, 3), dtype=np.uint8)
+mask[100:350, 350: -300] = 255
+mask = Image.fromarray(mask)
+mask.save("mask.jpg")
+
+image_2 = pipe(
+    prompt="a cat sitting on a chair, wearing sunglasses",
+    controlnet_inputs=[ControlNetInput(image=image_1, inpaint_mask=mask, scale=0.9)],
+    height=1024, width=1024,
+    seed=9, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference/FLUX.1-dev-Controlnet-Union-alpha.py

@@ -0,0 +1,40 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+from diffsynth.utils.controlnet import Annotator
+from modelscope import snapshot_download
+
+
+
+snapshot_download("sd_lora/Annotators", allow_file_pattern="dpt_hybrid-midas-501f0c75.pt", local_dir="models/Annotators")
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+image_1 = pipe(
+    prompt="a beautiful Asian girl, full body, red dress, summer",
+    height=1024, width=1024,
+    seed=6, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+image_canny = Annotator("canny")(image_1)
+image_depth = Annotator("depth")(image_1)
+
+image_2 = pipe(
+    prompt="a beautiful Asian girl, full body, red dress, winter",
+    controlnet_inputs=[
+        ControlNetInput(image=image_canny, scale=0.3, processor_id="canny"),
+        ControlNetInput(image=image_depth, scale=0.3, processor_id="depth"),
+    ],
+    height=1024, width=1024,
+    seed=7, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference/FLUX.1-dev-Controlnet-Upscaler.py

@@ -0,0 +1,33 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+image_1 = pipe(
+    prompt="a photo of a cat, highly detailed",
+    height=768, width=768,
+    seed=0, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+image_1 = image_1.resize((2048, 2048))
+image_2 = pipe(
+    prompt="a photo of a cat, highly detailed",
+    controlnet_inputs=[ControlNetInput(image=image_1, scale=0.7)],
+    input_image=image_1,
+    denoising_strength=0.99,
+    height=2048, width=2048, tiled=True,
+    seed=1, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference/FLUX.1-dev-EliGen.py

@@ -0,0 +1,133 @@
+import random
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+
+def visualize_masks(image, masks, mask_prompts, output_path, font_size=35, use_random_colors=False):
+    # Create a blank image for overlays
+    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
+    
+    colors = [
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+    ]
+    # Generate random colors for each mask
+    if use_random_colors:
+        colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 80) for _ in range(len(masks))]
+    
+    # Font settings
+    try:
+        font = ImageFont.truetype("arial", font_size)  # Adjust as needed
+    except IOError:
+        font = ImageFont.load_default(font_size)
+
+    # Overlay each mask onto the overlay image
+    for mask, mask_prompt, color in zip(masks, mask_prompts, colors):
+        # Convert mask to RGBA mode
+        mask_rgba = mask.convert('RGBA')
+        mask_data = mask_rgba.getdata()
+        new_data = [(color if item[:3] == (255, 255, 255) else (0, 0, 0, 0)) for item in mask_data]
+        mask_rgba.putdata(new_data)
+
+        # Draw the mask prompt text on the mask
+        draw = ImageDraw.Draw(mask_rgba)
+        mask_bbox = mask.getbbox()  # Get the bounding box of the mask
+        text_position = (mask_bbox[0] + 10, mask_bbox[1] + 10)  # Adjust text position based on mask position
+        draw.text(text_position, mask_prompt, fill=(255, 255, 255, 255), font=font)
+
+        # Alpha composite the overlay with this mask
+        overlay = Image.alpha_composite(overlay, mask_rgba)
+    
+    # Composite the overlay onto the original image
+    result = Image.alpha_composite(image.convert('RGBA'), overlay)
+    
+    # Save or display the resulting image
+    result.save(output_path)
+
+    return result
+
+def example(pipe, seeds, example_id, global_prompt, entity_prompts):
+    dataset_snapshot_download(dataset_id="DiffSynth-Studio/examples_in_diffsynth", local_dir="./", allow_file_pattern=f"data/examples/eligen/entity_control/example_{example_id}/*.png")
+    masks = [Image.open(f"./data/examples/eligen/entity_control/example_{example_id}/{i}.png").convert('RGB') for i in range(len(entity_prompts))]
+    negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+    for seed in seeds:
+        # generate image
+        image = pipe(
+            prompt=global_prompt,
+            cfg_scale=3.0,
+            negative_prompt=negative_prompt,
+            num_inference_steps=50,
+            embedded_guidance=3.5,
+            seed=seed,
+            height=1024,
+            width=1024,
+            eligen_entity_prompts=entity_prompts,
+            eligen_entity_masks=masks,
+        )
+        image.save(f"eligen_example_{example_id}_{seed}.png")
+        visualize_masks(image, masks, entity_prompts, f"eligen_example_{example_id}_mask_{seed}.png")
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+pipe.load_lora(pipe.dit, ModelConfig(model_id="DiffSynth-Studio/Eligen", origin_file_pattern="model_bf16.safetensors"), alpha=1)
+
+# example 1
+global_prompt = "A breathtaking beauty of Raja Ampat by the late-night moonlight , one beautiful woman from behind wearing a pale blue long dress with soft glow, sitting at the top of a cliff looking towards the beach,pastell light colors, a group of small distant birds flying in far sky, a boat sailing on the sea, best quality, realistic, whimsical, fantastic, splash art, intricate detailed, hyperdetailed, maximalist style, photorealistic, concept art, sharp focus, harmony, serenity, tranquility, soft pastell colors,ambient occlusion, cozy ambient lighting, masterpiece, liiv1, linquivera, metix, mentixis, masterpiece, award winning, view from above\n"
+entity_prompts = ["cliff", "sea", "moon", "sailing boat", "a seated beautiful woman", "pale blue long dress with soft glow"]
+example(pipe, [0], 1, global_prompt, entity_prompts)
+
+# example 2
+global_prompt = "samurai girl wearing a kimono, she's holding a sword  glowing with red flame, her long hair is flowing in the wind, she is looking at a small bird perched on the back of her hand. ultra realist style. maximum image detail. maximum realistic render."
+entity_prompts = ["flowing hair", "sword glowing with red flame", "A cute bird", "blue belt"]
+example(pipe, [0], 2, global_prompt, entity_prompts)
+
+# example 3
+global_prompt = "Image of a neverending staircase up to a mysterious palace in the sky, The ancient palace stood majestically atop a mist-shrouded mountain, sunrise, two traditional monk walk in the stair looking at the sunrise, fog,see-through, best quality, whimsical, fantastic, splash art, intricate detailed, hyperdetailed, photorealistic, concept art, harmony, serenity, tranquility, ambient occlusion, halation, cozy ambient lighting, dynamic lighting,masterpiece, liiv1, linquivera, metix, mentixis, masterpiece, award winning,"
+entity_prompts = ["ancient palace", "stone staircase with railings", "a traditional monk", "a traditional monk"]
+example(pipe, [27], 3, global_prompt, entity_prompts)
+
+# example 4
+global_prompt = "A beautiful girl wearing shirt and shorts in the street,  holding a sign 'Entity Control'"
+entity_prompts = ["A beautiful girl", "sign 'Entity Control'", "shorts", "shirt"]
+example(pipe, [21], 4, global_prompt, entity_prompts)
+
+# example 5
+global_prompt = "A captivating, dramatic scene in a painting that exudes mystery and foreboding. A white sky, swirling blue clouds, and a crescent yellow moon illuminate a solitary woman standing near the water's edge. Her long dress flows in the wind, silhouetted against the eerie glow. The water mirrors the fiery sky and moonlight, amplifying the uneasy atmosphere."
+entity_prompts = ["crescent yellow moon", "a solitary woman", "water", "swirling blue clouds"]
+example(pipe, [0], 5, global_prompt, entity_prompts)
+
+# example 6
+global_prompt = "Snow White and the 6 Dwarfs."
+entity_prompts = ["Dwarf 1", "Dwarf 2", "Dwarf 3", "Snow White", "Dwarf 4", "Dwarf 5", "Dwarf 6"]
+example(pipe, [8], 6, global_prompt, entity_prompts)
+
+# example 7, same prompt with different seeds
+seeds = range(5, 9)
+global_prompt = "A beautiful woman wearing white dress, holding a mirror, with a warm light background;"
+entity_prompts = ["A beautiful woman", "mirror", "necklace", "glasses", "earring", "white dress", "jewelry headpiece"]
+example(pipe, seeds, 7, global_prompt, entity_prompts)

examples/flux/model_inference/FLUX.1-dev-IP-Adapter.py

@@ -0,0 +1,24 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="InstantX/FLUX.1-dev-IP-Adapter", origin_file_pattern="ip-adapter.bin"),
+        ModelConfig(model_id="google/siglip-so400m-patch14-384", origin_file_pattern="model.safetensors"),
+    ],
+)
+
+origin_prompt = "a rabbit in a garden, colorful flowers"
+image = pipe(prompt=origin_prompt, height=1280, width=960, seed=42)
+image.save("style image.jpg")
+
+image = pipe(prompt="A piggy", height=1280, width=960, seed=42,
+    ipadapter_images=[image], ipadapter_scale=0.7)
+image.save("A piggy.jpg")

examples/flux/model_inference/FLUX.1-dev-InfiniteYou.py

@@ -0,0 +1,61 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+from modelscope import dataset_snapshot_download
+from modelscope import snapshot_download
+from PIL import Image
+import numpy as np
+
+# This model has additional requirements.
+# Please install the following packages.
+# pip install facexlib insightface onnxruntime
+snapshot_download(
+    "ByteDance/InfiniteYou",
+    allow_file_pattern="supports/insightface/models/antelopev2/*",
+    local_dir="models/ByteDance/InfiniteYou",
+)
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/image_proj_model.bin"),
+        ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors"),
+    ],
+)
+
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/examples_in_diffsynth",
+    local_dir="./",
+    allow_file_pattern=f"data/examples/infiniteyou/*",
+)
+
+height, width = 1024, 1024
+controlnet_image = Image.fromarray(np.zeros([height, width, 3]).astype(np.uint8))
+controlnet_inputs = [ControlNetInput(image=controlnet_image, scale=1.0, processor_id="None")]
+
+prompt = "A man, portrait, cinematic"
+id_image = "data/examples/infiniteyou/man.jpg"
+id_image = Image.open(id_image).convert('RGB')
+image = pipe(
+    prompt=prompt, seed=1,
+    infinityou_id_image=id_image, infinityou_guidance=1.0,
+    controlnet_inputs=controlnet_inputs,
+    num_inference_steps=50, embedded_guidance=3.5,
+    height=height, width=width,
+)
+image.save("man.jpg")
+
+prompt = "A woman, portrait, cinematic"
+id_image = "data/examples/infiniteyou/woman.jpg"
+id_image = Image.open(id_image).convert('RGB')
+image = pipe(
+    prompt=prompt, seed=1,
+    infinityou_id_image=id_image, infinityou_guidance=1.0,
+    controlnet_inputs=controlnet_inputs,
+    num_inference_steps=50, embedded_guidance=3.5,
+    height=height, width=width,
+)
+image.save("woman.jpg")

examples/flux/model_inference/FLUX.1-dev-LoRA-Encoder.py

@@ -0,0 +1,38 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev", origin_file_pattern="model.safetensors"),
+    ],
+)
+lora = ModelConfig(model_id="VoidOc/flux_animal_forest1", origin_file_pattern="20.safetensors")
+pipe.load_lora(pipe.dit, lora) # Use `pipe.clear_lora()` to drop the loaded LoRA.
+
+# Empty prompt can automatically activate LoRA capabilities.
+image = pipe(prompt="", seed=0, lora_encoder_inputs=lora)
+image.save("image_1.jpg")
+
+image = pipe(prompt="", seed=0)
+image.save("image_1_origin.jpg")
+
+# Prompt without trigger words can also activate LoRA capabilities.
+image = pipe(prompt="a car", seed=0, lora_encoder_inputs=lora)
+image.save("image_2.jpg")
+
+image = pipe(prompt="a car", seed=0,)
+image.save("image_2_origin.jpg")
+
+# Adjust the activation intensity through the scale parameter.
+image = pipe(prompt="a cat", seed=0, lora_encoder_inputs=lora, lora_encoder_scale=1.0)
+image.save("image_3.jpg")
+
+image = pipe(prompt="a cat", seed=0, lora_encoder_inputs=lora, lora_encoder_scale=0.5)
+image.save("image_3_scale.jpg")

examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py

@@ -0,0 +1,38 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+vram_config = {
+    # Enable lora hotloading
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cuda",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cuda",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/LoRAFusion-preview-FLUX.1-dev", origin_file_pattern="model.safetensors"),
+    ],
+)
+pipe.enable_lora_merger()
+
+pipe.load_lora(
+    pipe.dit,
+    ModelConfig(model_id="cancel13/cxsk", origin_file_pattern="30.safetensors"),
+)
+pipe.load_lora(
+    pipe.dit,
+    ModelConfig(model_id="DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1", origin_file_pattern="merged_lora.safetensors"),
+)
+image = pipe(prompt="a cat", seed=0)
+image.save("image_fused.jpg")

examples/flux/model_inference/FLUX.1-dev-Timestep-LoRA.py

@@ -0,0 +1,21 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+pipe.load_timestep_lora(ModelConfig(model_id="DiffSynth-Studio/MultiAlign-FLUX.1-dev", origin_file_pattern="adapter_model.safetensors"))
+
+prompt = "A moonlit Venetian canal scene with a gondola floating on the left, a candlelit table for two on a stone terrace to the right, and flower boxes under arched windows in the peach-colored building behind, with shimmering reflections in dark water; romantic cinematic mood and the restaurant sign reading 'NOTTE SERENA'"
+
+image = pipe(prompt=prompt, seed=0, num_inference_steps=30, t5_sequence_length=128)
+image.save("FLUX.1-dev-Timestep-LoRA.jpg")

examples/flux/model_inference/FLUX.1-dev.py

@@ -0,0 +1,26 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+)
+
+prompt = "CG, masterpiece, best quality, solo, long hair, wavy hair, silver hair, blue eyes, blue dress, medium breasts, dress, underwater, air bubble, floating hair, refraction, portrait. The girl's flowing silver hair shimmers with every color of the rainbow and cascades down, merging with the floating flora around her."
+negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+
+image = pipe(prompt=prompt, seed=0)
+image.save("flux.jpg")
+
+image = pipe(
+    prompt=prompt, negative_prompt=negative_prompt,
+    seed=0, cfg_scale=2, num_inference_steps=50,
+)
+image.save("flux_cfg.jpg")

examples/flux/model_inference/Insert-Anything.py

@@ -0,0 +1,44 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="flux1-fill-dev.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="ae.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_embedder/diffusion_pytorch_model.safetensors"),
+    ],
+)
+
+pipe.load_lora(pipe.dit, ModelConfig(model_id="HuanJue/Insert-Anything", origin_file_pattern="20250321_steps5000_pytorch_lora_weights.safetensors"))
+
+dataset_snapshot_download(
+    dataset_id="HuanJue/example_dataset",
+    local_dir="./",
+    allow_file_pattern=f"Insert-Anything/*",
+)
+
+source_image = Image.open("Insert-Anything/source_image.png").convert("RGB")
+source_mask = Image.open("Insert-Anything/source_mask.png").convert("L")
+ref_image = Image.open("Insert-Anything/ref_image.png").convert("RGB")
+ref_mask = Image.open("Insert-Anything/ref_mask.png").convert("L")
+
+seed = 666
+
+image = pipe(
+    insert_anything_source_image=source_image,
+    insert_anything_source_mask=source_mask,
+    insert_anything_ref_image=ref_image,
+    insert_anything_ref_mask=ref_mask,
+    seed=seed,
+    embedded_guidance=30.0,
+    num_inference_steps=50,
+)
+
+image.save("image_Insert-Anything.jpg")

examples/flux/model_inference/Nexus-Gen-Editing.py

@@ -0,0 +1,37 @@
+import importlib
+import torch
+from PIL import Image
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+
+if importlib.util.find_spec("transformers") is None:
+    raise ImportError("You are using Nexus-GenV2. It depends on transformers, which is not installed. Please install it with `pip install transformers==4.49.0`.")
+else:
+    import transformers
+    assert transformers.__version__ == "4.49.0", "Nexus-GenV2 requires transformers==4.49.0, please install it with `pip install transformers==4.49.0`."
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+    nexus_gen_processor_config=ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="processor/"),
+)
+
+dataset_snapshot_download(dataset_id="DiffSynth-Studio/examples_in_diffsynth", local_dir="./", allow_file_pattern=f"data/examples/nexusgen/cat.jpg")
+ref_image = Image.open("data/examples/nexusgen/cat.jpg").convert("RGB")
+prompt = "Add a crown."
+image = pipe(
+    prompt=prompt, negative_prompt="",
+    seed=42, cfg_scale=2.0, num_inference_steps=50,
+    nexus_gen_reference_image=ref_image,
+    height=512, width=512,
+)
+image.save("cat_crown.jpg")

examples/flux/model_inference/Nexus-Gen-Generation.py

@@ -0,0 +1,32 @@
+import importlib
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+if importlib.util.find_spec("transformers") is None:
+    raise ImportError("You are using Nexus-GenV2. It depends on transformers, which is not installed. Please install it with `pip install transformers==4.49.0`.")
+else:
+    import transformers
+    assert transformers.__version__ == "4.49.0", "Nexus-GenV2 requires transformers==4.49.0, please install it with `pip install transformers==4.49.0`."
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors"),
+        ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"),
+    ],
+    nexus_gen_processor_config=ModelConfig("DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="processor"),
+)
+
+prompt = "一只可爱的猫咪"
+image = pipe(
+    prompt=prompt, negative_prompt="",
+    seed=0, cfg_scale=3, num_inference_steps=50,
+    height=1024, width=1024,
+)
+image.save("cat.jpg")

examples/flux/model_inference/Step1X-Edit.py

@@ -0,0 +1,32 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+import numpy as np
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Qwen/Qwen2.5-VL-7B-Instruct", origin_file_pattern="model-*.safetensors"),
+        ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors"),
+        ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="vae.safetensors"),
+    ],
+)
+
+image = Image.fromarray(np.zeros((1248, 832, 3), dtype=np.uint8) + 255)
+image = pipe(
+    prompt="draw red flowers in Chinese ink painting style",
+    step1x_reference_image=image,
+    width=832, height=1248, cfg_scale=6,
+    seed=1, rand_device='cuda'
+)
+image.save("image_1.jpg")
+
+image = pipe(
+    prompt="add more flowers in Chinese ink painting style",
+    step1x_reference_image=image,
+    width=832, height=1248, cfg_scale=6,
+    seed=2, rand_device='cuda'
+)
+image.save("image_2.jpg")

examples/flux/model_inference_low_vram/FLEX.2-preview.py

@@ -0,0 +1,61 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from diffsynth.utils.controlnet import Annotator
+import numpy as np
+from PIL import Image
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="ostris/Flex.2-preview", origin_file_pattern="Flex.2-preview.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl, long hair, red t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    seed=0
+)
+image.save("image_1.jpg")
+
+mask = np.zeros((1024, 1024, 3), dtype=np.uint8)
+mask[200:400, 400:700] = 255
+mask = Image.fromarray(mask)
+mask.save("image_mask.jpg")
+
+inpaint_image = image
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl with sunglasses, long hair, red t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    flex_inpaint_image=inpaint_image, flex_inpaint_mask=mask,
+    seed=4
+)
+image.save("image_2.jpg")
+
+control_image = Annotator("canny")(image)
+control_image.save("image_control.jpg")
+
+image = pipe(
+    prompt="portrait of a beautiful Asian girl with sunglasses, long hair, yellow t-shirt, sunshine, beach",
+    num_inference_steps=50, embedded_guidance=3.5,
+    flex_control_image=control_image,
+    seed=4
+)
+image.save("image_3.jpg")

examples/flux/model_inference_low_vram/FLUX.1-Fill-dev.py

@@ -0,0 +1,39 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="flux1-fill-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Fill-dev", origin_file_pattern="ae.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+dataset_snapshot_download(
+    dataset_id="HuanJue/example_dataset",
+    local_dir="./",
+    allow_file_pattern=f"FLUX.1-Fill-dev/*",
+)
+
+flux_fill_image = Image.open("FLUX.1-Fill-dev/cup.png").convert("RGB")
+flux_fill_mask = Image.open("FLUX.1-Fill-dev/cup_mask.png").convert("L")
+prompt = "a white paper cup"
+image = pipe(prompt=prompt, flux_fill_image=flux_fill_image, flux_fill_mask=flux_fill_mask, height=1632, width=1232, seed=0, embedded_guidance=30.0, num_inference_steps=50)
+
+image.save("image_FLUX.1-Fill-dev.jpg")

examples/flux/model_inference_low_vram/FLUX.1-Kontext-dev.py

@@ -0,0 +1,65 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Kontext-dev", origin_file_pattern="flux1-kontext-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image_1 = pipe(
+    prompt="a beautiful Asian long-haired female college student.",
+    embedded_guidance=2.5,
+    seed=1,
+)
+image_1.save("image_1.jpg")
+
+image_2 = pipe(
+    prompt="transform the style to anime style.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=2,
+)
+image_2.save("image_2.jpg")
+
+image_3 = pipe(
+    prompt="let her smile.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=3,
+)
+image_3.save("image_3.jpg")
+
+image_4 = pipe(
+    prompt="let the girl play basketball.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=4,
+)
+image_4.save("image_4.jpg")
+
+image_5 = pipe(
+    prompt="move the girl to a park, let her sit on a chair.",
+    kontext_images=image_1,
+    embedded_guidance=2.5,
+    seed=5,
+)
+image_5.save("image_5.jpg")

examples/flux/model_inference_low_vram/FLUX.1-Krea-dev.py

@@ -0,0 +1,38 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Krea-dev", origin_file_pattern="flux1-krea-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+prompt = "An beautiful woman is riding a bicycle in a park, wearing a red dress"
+negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+
+image = pipe(prompt=prompt, seed=0, embedded_guidance=4.5)
+image.save("flux_krea.jpg")
+
+image = pipe(
+    prompt=prompt, negative_prompt=negative_prompt,
+    seed=0, cfg_scale=2, num_inference_steps=50,
+    embedded_guidance=4.5
+)
+image.save("flux_krea_cfg.jpg")

examples/flux/model_inference_low_vram/FLUX.1-Redux-dev.py

@@ -0,0 +1,40 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from PIL import Image
+from modelscope import dataset_snapshot_download
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-Redux-dev", origin_file_pattern="image_embedder/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+dataset_snapshot_download(
+    dataset_id="HuanJue/example_dataset",
+    local_dir="./",
+    allow_file_pattern=f"FLUX.1-Redux-dev/*",
+)
+
+flux_redux_image = Image.open("FLUX.1-Redux-dev/robot.png").convert("RGB")
+
+image = pipe(flux_redux_image=flux_redux_image, embedded_guidance=2.5, num_inference_steps=50)
+
+image.save("image_FLUX.1-Redux-dev.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py

@@ -0,0 +1,30 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors", **vram_config)
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+for i in [0.1, 0.3, 0.5, 0.7, 0.9]:
+    image = pipe(prompt="a cat on the beach", seed=2, value_controller_inputs=[i])
+    image.save(f"value_control_{i}.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Inpainting-Beta.py

@@ -0,0 +1,48 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+import numpy as np
+from PIL import Image
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", origin_file_pattern="diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image_1 = pipe(
+    prompt="a cat sitting on a chair",
+    height=1024, width=1024,
+    seed=8, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+mask = np.zeros((1024, 1024, 3), dtype=np.uint8)
+mask[100:350, 350: -300] = 255
+mask = Image.fromarray(mask)
+mask.save("mask.jpg")
+
+image_2 = pipe(
+    prompt="a cat sitting on a chair, wearing sunglasses",
+    controlnet_inputs=[ControlNetInput(image=image_1, inpaint_mask=mask, scale=0.9)],
+    height=1024, width=1024,
+    seed=9, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Union-alpha.py

@@ -0,0 +1,50 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+from diffsynth.utils.controlnet import Annotator
+from modelscope import snapshot_download
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+snapshot_download("sd_lora/Annotators", allow_file_pattern="dpt_hybrid-midas-501f0c75.pt", local_dir="models/Annotators")
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image_1 = pipe(
+    prompt="a beautiful Asian girl, full body, red dress, summer",
+    height=1024, width=1024,
+    seed=6, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+image_canny = Annotator("canny")(image_1)
+image_depth = Annotator("depth")(image_1)
+
+image_2 = pipe(
+    prompt="a beautiful Asian girl, full body, red dress, winter",
+    controlnet_inputs=[
+        ControlNetInput(image=image_canny, scale=0.3, processor_id="canny"),
+        ControlNetInput(image=image_depth, scale=0.3, processor_id="depth"),
+    ],
+    height=1024, width=1024,
+    seed=7, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Upscaler.py

@@ -0,0 +1,44 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+image_1 = pipe(
+    prompt="a photo of a cat, highly detailed",
+    height=768, width=768,
+    seed=0, rand_device="cuda",
+)
+image_1.save("image_1.jpg")
+
+image_1 = image_1.resize((2048, 2048))
+image_2 = pipe(
+    prompt="a photo of a cat, highly detailed",
+    controlnet_inputs=[ControlNetInput(image=image_1, scale=0.7)],
+    input_image=image_1,
+    denoising_strength=0.99,
+    height=2048, width=2048, tiled=True,
+    seed=1, rand_device="cuda",
+)
+image_2.save("image_2.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-EliGen.py

@@ -0,0 +1,144 @@
+import random
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+from modelscope import dataset_snapshot_download
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+def visualize_masks(image, masks, mask_prompts, output_path, font_size=35, use_random_colors=False):
+    # Create a blank image for overlays
+    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
+    
+    colors = [
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+        (165, 238, 173, 80),
+        (76, 102, 221, 80),
+        (221, 160, 77, 80),
+        (204, 93, 71, 80),
+        (145, 187, 149, 80),
+        (134, 141, 172, 80),
+        (157, 137, 109, 80),
+        (153, 104, 95, 80),
+    ]
+    # Generate random colors for each mask
+    if use_random_colors:
+        colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 80) for _ in range(len(masks))]
+    
+    # Font settings
+    try:
+        font = ImageFont.truetype("arial", font_size)  # Adjust as needed
+    except IOError:
+        font = ImageFont.load_default(font_size)
+
+    # Overlay each mask onto the overlay image
+    for mask, mask_prompt, color in zip(masks, mask_prompts, colors):
+        # Convert mask to RGBA mode
+        mask_rgba = mask.convert('RGBA')
+        mask_data = mask_rgba.getdata()
+        new_data = [(color if item[:3] == (255, 255, 255) else (0, 0, 0, 0)) for item in mask_data]
+        mask_rgba.putdata(new_data)
+
+        # Draw the mask prompt text on the mask
+        draw = ImageDraw.Draw(mask_rgba)
+        mask_bbox = mask.getbbox()  # Get the bounding box of the mask
+        text_position = (mask_bbox[0] + 10, mask_bbox[1] + 10)  # Adjust text position based on mask position
+        draw.text(text_position, mask_prompt, fill=(255, 255, 255, 255), font=font)
+
+        # Alpha composite the overlay with this mask
+        overlay = Image.alpha_composite(overlay, mask_rgba)
+    
+    # Composite the overlay onto the original image
+    result = Image.alpha_composite(image.convert('RGBA'), overlay)
+    
+    # Save or display the resulting image
+    result.save(output_path)
+
+    return result
+
+def example(pipe, seeds, example_id, global_prompt, entity_prompts):
+    dataset_snapshot_download(dataset_id="DiffSynth-Studio/examples_in_diffsynth", local_dir="./", allow_file_pattern=f"data/examples/eligen/entity_control/example_{example_id}/*.png")
+    masks = [Image.open(f"./data/examples/eligen/entity_control/example_{example_id}/{i}.png").convert('RGB') for i in range(len(entity_prompts))]
+    negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
+    for seed in seeds:
+        # generate image
+        image = pipe(
+            prompt=global_prompt,
+            cfg_scale=3.0,
+            negative_prompt=negative_prompt,
+            num_inference_steps=50,
+            embedded_guidance=3.5,
+            seed=seed,
+            height=1024,
+            width=1024,
+            eligen_entity_prompts=entity_prompts,
+            eligen_entity_masks=masks,
+        )
+        image.save(f"eligen_example_{example_id}_{seed}.png")
+        visualize_masks(image, masks, entity_prompts, f"eligen_example_{example_id}_mask_{seed}.png")
+
+
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+pipe.load_lora(pipe.dit, ModelConfig(model_id="DiffSynth-Studio/Eligen", origin_file_pattern="model_bf16.safetensors"), alpha=1)
+
+# example 1
+global_prompt = "A breathtaking beauty of Raja Ampat by the late-night moonlight , one beautiful woman from behind wearing a pale blue long dress with soft glow, sitting at the top of a cliff looking towards the beach,pastell light colors, a group of small distant birds flying in far sky, a boat sailing on the sea, best quality, realistic, whimsical, fantastic, splash art, intricate detailed, hyperdetailed, maximalist style, photorealistic, concept art, sharp focus, harmony, serenity, tranquility, soft pastell colors,ambient occlusion, cozy ambient lighting, masterpiece, liiv1, linquivera, metix, mentixis, masterpiece, award winning, view from above\n"
+entity_prompts = ["cliff", "sea", "moon", "sailing boat", "a seated beautiful woman", "pale blue long dress with soft glow"]
+example(pipe, [0], 1, global_prompt, entity_prompts)
+
+# example 2
+global_prompt = "samurai girl wearing a kimono, she's holding a sword  glowing with red flame, her long hair is flowing in the wind, she is looking at a small bird perched on the back of her hand. ultra realist style. maximum image detail. maximum realistic render."
+entity_prompts = ["flowing hair", "sword glowing with red flame", "A cute bird", "blue belt"]
+example(pipe, [0], 2, global_prompt, entity_prompts)
+
+# example 3
+global_prompt = "Image of a neverending staircase up to a mysterious palace in the sky, The ancient palace stood majestically atop a mist-shrouded mountain, sunrise, two traditional monk walk in the stair looking at the sunrise, fog,see-through, best quality, whimsical, fantastic, splash art, intricate detailed, hyperdetailed, photorealistic, concept art, harmony, serenity, tranquility, ambient occlusion, halation, cozy ambient lighting, dynamic lighting,masterpiece, liiv1, linquivera, metix, mentixis, masterpiece, award winning,"
+entity_prompts = ["ancient palace", "stone staircase with railings", "a traditional monk", "a traditional monk"]
+example(pipe, [27], 3, global_prompt, entity_prompts)
+
+# example 4
+global_prompt = "A beautiful girl wearing shirt and shorts in the street,  holding a sign 'Entity Control'"
+entity_prompts = ["A beautiful girl", "sign 'Entity Control'", "shorts", "shirt"]
+example(pipe, [21], 4, global_prompt, entity_prompts)
+
+# example 5
+global_prompt = "A captivating, dramatic scene in a painting that exudes mystery and foreboding. A white sky, swirling blue clouds, and a crescent yellow moon illuminate a solitary woman standing near the water's edge. Her long dress flows in the wind, silhouetted against the eerie glow. The water mirrors the fiery sky and moonlight, amplifying the uneasy atmosphere."
+entity_prompts = ["crescent yellow moon", "a solitary woman", "water", "swirling blue clouds"]
+example(pipe, [0], 5, global_prompt, entity_prompts)
+
+# example 6
+global_prompt = "Snow White and the 6 Dwarfs."
+entity_prompts = ["Dwarf 1", "Dwarf 2", "Dwarf 3", "Snow White", "Dwarf 4", "Dwarf 5", "Dwarf 6"]
+example(pipe, [8], 6, global_prompt, entity_prompts)
+
+# example 7, same prompt with different seeds
+seeds = range(5, 9)
+global_prompt = "A beautiful woman wearing white dress, holding a mirror, with a warm light background;"
+entity_prompts = ["A beautiful woman", "mirror", "necklace", "glasses", "earring", "white dress", "jewelry headpiece"]
+example(pipe, seeds, 7, global_prompt, entity_prompts)

examples/flux/model_inference_low_vram/FLUX.1-dev-IP-Adapter.py

@@ -0,0 +1,35 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig
+
+
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="InstantX/FLUX.1-dev-IP-Adapter", origin_file_pattern="ip-adapter.bin", **vram_config),
+        ModelConfig(model_id="google/siglip-so400m-patch14-384", origin_file_pattern="model.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+origin_prompt = "a rabbit in a garden, colorful flowers"
+image = pipe(prompt=origin_prompt, height=1280, width=960, seed=42)
+image.save("style image.jpg")
+
+image = pipe(prompt="A piggy", height=1280, width=960, seed=42,
+    ipadapter_images=[image], ipadapter_scale=0.7)
+image.save("A piggy.jpg")

examples/flux/model_inference_low_vram/FLUX.1-dev-InfiniteYou.py

@@ -0,0 +1,73 @@
+import torch
+from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput
+from modelscope import dataset_snapshot_download
+from modelscope import snapshot_download
+from PIL import Image
+import numpy as np
+
+
+# This model has additional requirements.
+# Please install the following packages.
+# pip install facexlib insightface onnxruntime
+vram_config = {
+    "offload_dtype": torch.float8_e4m3fn,
+    "offload_device": "cpu",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+snapshot_download(
+    "ByteDance/InfiniteYou",
+    allow_file_pattern="supports/insightface/models/antelopev2/*",
+    local_dir="models/ByteDance/InfiniteYou",
+)
+pipe = FluxImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/*.safetensors", **vram_config),
+        ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", **vram_config),
+        ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/image_proj_model.bin", **vram_config),
+        ModelConfig(model_id="ByteDance/InfiniteYou", origin_file_pattern="infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors", **vram_config),
+    ],
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
+)
+
+dataset_snapshot_download(
+    dataset_id="DiffSynth-Studio/examples_in_diffsynth",
+    local_dir="./",
+    allow_file_pattern=f"data/examples/infiniteyou/*",
+)
+
+height, width = 1024, 1024
+controlnet_image = Image.fromarray(np.zeros([height, width, 3]).astype(np.uint8))
+controlnet_inputs = [ControlNetInput(image=controlnet_image, scale=1.0, processor_id="None")]
+
+prompt = "A man, portrait, cinematic"
+id_image = "data/examples/infiniteyou/man.jpg"
+id_image = Image.open(id_image).convert('RGB')
+image = pipe(
+    prompt=prompt, seed=1,
+    infinityou_id_image=id_image, infinityou_guidance=1.0,
+    controlnet_inputs=controlnet_inputs,
+    num_inference_steps=50, embedded_guidance=3.5,
+    height=height, width=width,
+)
+image.save("man.jpg")
+
+prompt = "A woman, portrait, cinematic"
+id_image = "data/examples/infiniteyou/woman.jpg"
+id_image = Image.open(id_image).convert('RGB')
+image = pipe(
+    prompt=prompt, seed=1,
+    infinityou_id_image=id_image, infinityou_guidance=1.0,
+    controlnet_inputs=controlnet_inputs,
+    num_inference_steps=50, embedded_guidance=3.5,
+    height=height, width=width,
+)
+image.save("woman.jpg")