diff --git a/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Edit.py b/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..98176be9e7bc3a7677eb76df43ebe29840b622c7 --- /dev/null +++ b/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") diff --git a/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Turbo.py b/examples/boogu_image/model_training/validate_lora/Boogu-Image-0.1-Turbo.py new file mode 100644 index 0000000000000000000000000000000000000000..4562ec6ddd9cc1ffebf4b04b2209c17ba14492d1 --- /dev/null +++ b/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") diff --git a/examples/dev_tools/unit_test.py b/examples/dev_tools/unit_test.py new file mode 100644 index 0000000000000000000000000000000000000000..200ced8e006ab6b4def8c4c924738c639ec9ca1f --- /dev/null +++ b/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() diff --git a/examples/dev_tools/webui.py b/examples/dev_tools/webui.py new file mode 100644 index 0000000000000000000000000000000000000000..24ed99c22b8f255b6d923fd5df2579078857c61a --- /dev/null +++ b/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() diff --git a/examples/dev_tools/webui_train.py b/examples/dev_tools/webui_train.py new file mode 100644 index 0000000000000000000000000000000000000000..0726233ef06824e6d3fe5f3cf0e832c6fc21ef3f --- /dev/null +++ b/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() diff --git a/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py b/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py new file mode 100644 index 0000000000000000000000000000000000000000..afec5e58296ab28b5f954ef8e44ace7722bab282 --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_inference/ERNIE-Image.py b/examples/ernie_image/model_inference/ERNIE-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..29e78e8b0d7f28b195ae418f223cb00b4de0e029 --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py b/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py new file mode 100644 index 0000000000000000000000000000000000000000..64b928c3f4b4d4ba0112ae60bd13f619f591aca2 --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py b/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..ca7494715496bbf8870abc7574955622c14ed233 --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_training/full/ERNIE-Image.sh b/examples/ernie_image/model_training/full/ERNIE-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..374c1a841bdd2f12d2568661086af1dc6c3cab42 --- /dev/null +++ b/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 diff --git a/examples/ernie_image/model_training/full/accelerate_config_zero3.yaml b/examples/ernie_image/model_training/full/accelerate_config_zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e6a8d273346c6e04e7ab97c8aa661914a0819a86 --- /dev/null +++ b/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 diff --git a/examples/ernie_image/model_training/lora/ERNIE-Image.sh b/examples/ernie_image/model_training/lora/ERNIE-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..20f4ad5c433de46b4f8421a290e1169afe82b560 --- /dev/null +++ b/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 diff --git a/examples/ernie_image/model_training/special/split_training/ERNIE-Image.sh b/examples/ernie_image/model_training/special/split_training/ERNIE-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..fab7aeb824c1b364451ada2c168ff7f6d2268ba5 --- /dev/null +++ b/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 diff --git a/examples/ernie_image/model_training/special/split_training/validate.py b/examples/ernie_image/model_training/special/split_training/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..d2694d7d2379e663ddcb1d02bf931ce5c3a9e8dd --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_training/train.py b/examples/ernie_image/model_training/train.py new file mode 100644 index 0000000000000000000000000000000000000000..f2ff8dd1d161f0fa79f2110d5eb03d1bc9479052 --- /dev/null +++ b/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) diff --git a/examples/ernie_image/model_training/validate_full/ERNIE-Image.py b/examples/ernie_image/model_training/validate_full/ERNIE-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..11532b31a1a731f9a8e929f370cf43638b7bc75b --- /dev/null +++ b/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") diff --git a/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py b/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..787b81135dc930d06ffa1871110700965d6086ba --- /dev/null +++ b/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") diff --git a/examples/flux/README.md b/examples/flux/README.md new file mode 100644 index 0000000000000000000000000000000000000000..147a59d984d45faf579f087e244ca484ead69650 --- /dev/null +++ b/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 diff --git a/examples/flux/model_inference/FLEX.2-preview.py b/examples/flux/model_inference/FLEX.2-preview.py new file mode 100644 index 0000000000000000000000000000000000000000..efc8e9143c299b3a408967dbde898f78f959a95a --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-Fill-dev.py b/examples/flux/model_inference/FLUX.1-Fill-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..46db1a6df2e06c86e66a34d966cda288f0bab11a --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-Kontext-dev.py b/examples/flux/model_inference/FLUX.1-Kontext-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..e7aae1bb2410d397585082607697aa5a778a64f0 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-Krea-dev.py b/examples/flux/model_inference/FLUX.1-Krea-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..978a26a90d460c1723b4865ff5e36e5d98cddd17 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-Redux-dev.py b/examples/flux/model_inference/FLUX.1-Redux-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..bc12d37ecb706b7f890bbd36ce3d79a684099c00 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000000000000000000000000000000000000..b35cce8768f3aae7dbc7b7a6ddf7fd50eb029d6d --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-dev-Controlnet-Inpainting-Beta.py b/examples/flux/model_inference/FLUX.1-dev-Controlnet-Inpainting-Beta.py new file mode 100644 index 0000000000000000000000000000000000000000..3a0d1f3a8259a61f3cd6cb1b5fbcf256f9f43c24 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-dev-Controlnet-Union-alpha.py b/examples/flux/model_inference/FLUX.1-dev-Controlnet-Union-alpha.py new file mode 100644 index 0000000000000000000000000000000000000000..2fa10aac083ad338028aeb30a99b63efdfa656d9 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-dev-Controlnet-Upscaler.py b/examples/flux/model_inference/FLUX.1-dev-Controlnet-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..b4c288df1668b8731e5a80f367347819408f71bf --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-dev-EliGen.py b/examples/flux/model_inference/FLUX.1-dev-EliGen.py new file mode 100644 index 0000000000000000000000000000000000000000..6bc4d2e1bbe3b5591666f15e00d9098d562f4b6b --- /dev/null +++ b/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) diff --git a/examples/flux/model_inference/FLUX.1-dev-IP-Adapter.py b/examples/flux/model_inference/FLUX.1-dev-IP-Adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..1479e1da4bf59a1ab0b669147cb3488bcff3c66b --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-dev-InfiniteYou.py b/examples/flux/model_inference/FLUX.1-dev-InfiniteYou.py new file mode 100644 index 0000000000000000000000000000000000000000..4491ccb48e8642b4b1afe8a9bce74df0059b9b79 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-dev-LoRA-Encoder.py b/examples/flux/model_inference/FLUX.1-dev-LoRA-Encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..75f1bc80444aebd54290b46be9e4d9eb3c8d2e8d --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py b/examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..5076a7abd13775f9ca235ca5b1f4773636419c8d --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/FLUX.1-dev-Timestep-LoRA.py b/examples/flux/model_inference/FLUX.1-dev-Timestep-LoRA.py new file mode 100644 index 0000000000000000000000000000000000000000..d85cf18591ebbec62660e33658b311a711d231c2 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference/FLUX.1-dev.py b/examples/flux/model_inference/FLUX.1-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..35d1e96123f6612722b59c0bc5ab31abd08774c4 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/Insert-Anything.py b/examples/flux/model_inference/Insert-Anything.py new file mode 100644 index 0000000000000000000000000000000000000000..ee1ecc21fcff8d073947fb836361decee4c95d6a --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/Nexus-Gen-Editing.py b/examples/flux/model_inference/Nexus-Gen-Editing.py new file mode 100644 index 0000000000000000000000000000000000000000..67691659a0bc375b77f99bc0fa60c467f22ba787 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/Nexus-Gen-Generation.py b/examples/flux/model_inference/Nexus-Gen-Generation.py new file mode 100644 index 0000000000000000000000000000000000000000..5130d670672b6db41c71abf900e966ab862be57e --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference/Step1X-Edit.py b/examples/flux/model_inference/Step1X-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec517b6c35763e8f0a332d44e93bf5dfb282784 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLEX.2-preview.py b/examples/flux/model_inference_low_vram/FLEX.2-preview.py new file mode 100644 index 0000000000000000000000000000000000000000..a4454e84fedaf6405ee2ae7db152a234edea4b98 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-Fill-dev.py b/examples/flux/model_inference_low_vram/FLUX.1-Fill-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..f3f00b1004cf099a001dd51f9faf1c787e4e1d67 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-Kontext-dev.py b/examples/flux/model_inference_low_vram/FLUX.1-Kontext-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..2994a3301c9c9d57d88a7cfcf43b2c8bdde54812 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference_low_vram/FLUX.1-Krea-dev.py b/examples/flux/model_inference_low_vram/FLUX.1-Krea-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..2ceb064f600c233e4061ee057bebb55184f964df --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-Redux-dev.py b/examples/flux/model_inference_low_vram/FLUX.1-Redux-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..ea3a63207f6034a7022b4750b95ca2521687839e --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000000000000000000000000000000000000..e0226ba79ffef81e629e884ae1d067d34647389e --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Inpainting-Beta.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Inpainting-Beta.py new file mode 100644 index 0000000000000000000000000000000000000000..61ac25f805a2fa0f0ddcf977ad9b507e29d00de2 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Union-alpha.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Union-alpha.py new file mode 100644 index 0000000000000000000000000000000000000000..148e7ef95f08933bf1c4128ee59c7a31aa418bb0 --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Upscaler.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-Controlnet-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..ca7c72c04379e1f7dbcb3b935f3f9255719c5b67 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-EliGen.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-EliGen.py new file mode 100644 index 0000000000000000000000000000000000000000..da0d7cacb442fa1d8c244ec56ccf7d319a9c8055 --- /dev/null +++ b/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) diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-IP-Adapter.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-IP-Adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..59f2e9f9006be9186d5181ab397fa48407b52f1a --- /dev/null +++ b/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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-InfiniteYou.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-InfiniteYou.py new file mode 100644 index 0000000000000000000000000000000000000000..119856afd98f2c2b088debc940bbcc1d692d1d9c --- /dev/null +++ b/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") \ No newline at end of file diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Encoder.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..5928af0edf7832edc5126618ce8e52c04fd9f6b6 --- /dev/null +++ b/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Encoder.py @@ -0,0 +1,49 @@ +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/LoRA-Encoder-FLUX.1-Dev", origin_file_pattern="model.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Fusion.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..ce587cdac6d12923050b01c6003133274ba5f7bd --- /dev/null +++ b/examples/flux/model_inference_low_vram/FLUX.1-dev-LoRA-Fusion.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-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/LoRAFusion-preview-FLUX.1-dev", origin_file_pattern="model.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +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") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev.py b/examples/flux/model_inference_low_vram/FLUX.1-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..ffaf1813fe53a19bdb21f44313ddbe388a124cb1 --- /dev/null +++ b/examples/flux/model_inference_low_vram/FLUX.1-dev.py @@ -0,0 +1,37 @@ +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), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +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") diff --git a/examples/flux/model_inference_low_vram/Insert-Anything.py b/examples/flux/model_inference_low_vram/Insert-Anything.py new file mode 100644 index 0000000000000000000000000000000000000000..76dac9902b7da40d81353dc1e6b3b13e6f7d2924 --- /dev/null +++ b/examples/flux/model_inference_low_vram/Insert-Anything.py @@ -0,0 +1,56 @@ +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), + 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, +) + +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") diff --git a/examples/flux/model_inference_low_vram/Nexus-Gen-Editing.py b/examples/flux/model_inference_low_vram/Nexus-Gen-Editing.py new file mode 100644 index 0000000000000000000000000000000000000000..1b3050fbfa92bff5ddf6a7a9da5c5ab3da67591a --- /dev/null +++ b/examples/flux/model_inference_low_vram/Nexus-Gen-Editing.py @@ -0,0 +1,48 @@ +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`." + + +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="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="edit_decoder.bin", **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), + ], + nexus_gen_processor_config=ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="processor/"), + 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/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") diff --git a/examples/flux/model_inference_low_vram/Nexus-Gen-Generation.py b/examples/flux/model_inference_low_vram/Nexus-Gen-Generation.py new file mode 100644 index 0000000000000000000000000000000000000000..8372fcb822d0f0862a4352a77d0c329f60128b5f --- /dev/null +++ b/examples/flux/model_inference_low_vram/Nexus-Gen-Generation.py @@ -0,0 +1,43 @@ +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`." + + +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="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="model*.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="generation_decoder.bin", **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), + ], + nexus_gen_processor_config=ModelConfig("DiffSynth-Studio/Nexus-GenV2", origin_file_pattern="processor"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "一只可爱的猫咪" +image = pipe( + prompt=prompt, negative_prompt="", + seed=0, cfg_scale=3, num_inference_steps=50, + height=1024, width=1024, +) +image.save("cat.jpg") diff --git a/examples/flux/model_inference_low_vram/Step1X-Edit.py b/examples/flux/model_inference_low_vram/Step1X-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..9a3bde88ae112eced27d9ca7d7d10743c1df803f --- /dev/null +++ b/examples/flux/model_inference_low_vram/Step1X-Edit.py @@ -0,0 +1,43 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from PIL import Image +import numpy as np + + +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="Qwen/Qwen2.5-VL-7B-Instruct", origin_file_pattern="model-*.safetensors", **vram_config), + ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="step1x-edit-i1258.safetensors", **vram_config), + ModelConfig(model_id="stepfun-ai/Step1X-Edit", origin_file_pattern="vae.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +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") diff --git a/examples/flux/model_training/full/FLEX.2-preview.sh b/examples/flux/model_training/full/FLEX.2-preview.sh new file mode 100644 index 0000000000000000000000000000000000000000..0e8116eb9379a126a0a9358269e275e063745191 --- /dev/null +++ b/examples/flux/model_training/full/FLEX.2-preview.sh @@ -0,0 +1,14 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLEX.2-preview/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLEX.2-preview \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLEX.2-preview/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 200 \ + --model_id_with_origin_paths "ostris/Flex.2-preview:Flex.2-preview.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLEX.2-preview_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-Fill-dev.sh b/examples/flux/model_training/full/FLUX.1-Fill-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..4de3c5f47a83382a2ee0850aa685275503c0f744 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-Fill-dev.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Fill-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/metadata.csv \ + --data_file_keys "image,flux_fill_image,flux_fill_mask" \ + --max_pixels 1048576 \ + --dataset_repeat 200 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Fill-dev:flux1-fill-dev.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-Fill-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Fill-dev_full" \ + --trainable_models "dit" \ + --extra_inputs "flux_fill_image,flux_fill_mask" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-Kontext-dev.sh b/examples/flux/model_training/full/FLUX.1-Kontext-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..dd7cff7a39d16c96e6c5a4b76d1a42f36801d698 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-Kontext-dev.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Kontext-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev/metadata.csv \ + --data_file_keys "image,kontext_images" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Kontext-dev:flux1-kontext-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Kontext-dev_full" \ + --trainable_models "dit" \ + --extra_inputs "kontext_images" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-Krea-dev.sh b/examples/flux/model_training/full/FLUX.1-Krea-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..d9d685a58895aaf7c0c402d62d52a085e9348850 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-Krea-dev.sh @@ -0,0 +1,14 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Krea-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Krea-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Krea-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Krea-dev:flux1-krea-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Krea-dev_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-Redux-dev.sh b/examples/flux/model_training/full/FLUX.1-Redux-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..f1f768faf8ffa3d65c9bc6e1832184c358f11510 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-Redux-dev.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Redux-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev/metadata.csv \ + --data_file_keys "image,flux_redux_image" \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_encoder/model.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_embedder/diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Redux-dev_full" \ + --trainable_models "dit" \ + --extra_inputs "flux_redux_image" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh b/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh new file mode 100644 index 0000000000000000000000000000000000000000..1825ac266ef6e5631d30bb73a8d42b29eb938c45 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-AttriCtrl/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-AttriCtrl \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-AttriCtrl/metadata.csv \ + --data_file_keys "image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,DiffSynth-Studio/AttriCtrl-FLUX.1-Dev:models/brightness.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.value_controller.encoders.0." \ + --output_path "./models/train/FLUX.1-dev-AttriCtrl_full" \ + --trainable_models "value_controller" \ + --extra_inputs "value_controller_inputs" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Inpainting-Beta.sh b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Inpainting-Beta.sh new file mode 100644 index 0000000000000000000000000000000000000000..65241564a3d81d1f30881bab106b4eb145e0c59e --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Inpainting-Beta.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Inpainting-Beta/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Inpainting-Beta \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Inpainting-Beta/metadata.csv \ + --data_file_keys "image,controlnet_image,controlnet_inpaint_mask" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.controlnet.models.0." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Inpainting-Beta_full" \ + --trainable_models "controlnet" \ + --extra_inputs "controlnet_image,controlnet_inpaint_mask" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Union-alpha.sh b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Union-alpha.sh new file mode 100644 index 0000000000000000000000000000000000000000..d5c9c54e72e8476738422ff340d2d7d90ae1f8f6 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Union-alpha.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Union-alpha/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Union-alpha \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Union-alpha/metadata.csv \ + --data_file_keys "image,controlnet_image" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,InstantX/FLUX.1-dev-Controlnet-Union-alpha:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.controlnet.models.0." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Union-alpha_full" \ + --trainable_models "controlnet" \ + --extra_inputs "controlnet_image,controlnet_processor_id" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Upscaler.sh b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Upscaler.sh new file mode 100644 index 0000000000000000000000000000000000000000..99f7e220e7a23bd36721662a9feb44454ba01f41 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-Controlnet-Upscaler.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Upscaler/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Upscaler \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Upscaler/metadata.csv \ + --data_file_keys "image,controlnet_image" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,jasperai/Flux.1-dev-Controlnet-Upscaler:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.controlnet.models.0." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Upscaler_full" \ + --trainable_models "controlnet" \ + --extra_inputs "controlnet_image" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-IP-Adapter.sh b/examples/flux/model_training/full/FLUX.1-dev-IP-Adapter.sh new file mode 100644 index 0000000000000000000000000000000000000000..5600dcbca9c8607f3cfc875a57704ab7a17ddacf --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-IP-Adapter.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-IP-Adapter/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-IP-Adapter \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-IP-Adapter/metadata.csv \ + --data_file_keys "image,ipadapter_images" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,InstantX/FLUX.1-dev-IP-Adapter:ip-adapter.bin,google/siglip-so400m-patch14-384:model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.ipadapter." \ + --output_path "./models/train/FLUX.1-dev-IP-Adapter_full" \ + --trainable_models "ipadapter" \ + --extra_inputs "ipadapter_images" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-InfiniteYou.sh b/examples/flux/model_training/full/FLUX.1-dev-InfiniteYou.sh new file mode 100644 index 0000000000000000000000000000000000000000..a7d6c46b40aa5d3a68eb08fd5cad1571bed89b82 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-InfiniteYou.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-InfiniteYou/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-InfiniteYou \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-InfiniteYou/metadata.csv \ + --data_file_keys "image,controlnet_image,infinityou_id_image" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,ByteDance/InfiniteYou:infu_flux_v1.0/aes_stage2/image_proj_model.bin,ByteDance/InfiniteYou:infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe." \ + --output_path "./models/train/FLUX.1-dev-InfiniteYou_full" \ + --trainable_models "controlnet,image_proj_model" \ + --extra_inputs "controlnet_image,infinityou_id_image,infinityou_guidance" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev-LoRA-Encoder.sh b/examples/flux/model_training/full/FLUX.1-dev-LoRA-Encoder.sh new file mode 100644 index 0000000000000000000000000000000000000000..2240fd62534b9f19a9439b4e34cc4c96f7e6483a --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-LoRA-Encoder.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-LoRA-Encoder/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-LoRA-Encoder \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-LoRA-Encoder/metadata.csv \ + --data_file_keys "image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,DiffSynth-Studio/LoRA-Encoder-FLUX.1-Dev:model.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.lora_encoder." \ + --output_path "./models/train/FLUX.1-dev-LoRA-Encoder_full" \ + --trainable_models "lora_encoder" \ + --extra_inputs "lora_encoder_inputs" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/FLUX.1-dev.sh b/examples/flux/model_training/full/FLUX.1-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..6f126a08ca71bc0f206b9608c938aeb18361aac0 --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev.sh @@ -0,0 +1,14 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/full/Nexus-Gen.sh b/examples/flux/model_training/full/Nexus-Gen.sh new file mode 100644 index 0000000000000000000000000000000000000000..2c8242c14e0643aa601a5396512343dd4ddd6f90 --- /dev/null +++ b/examples/flux/model_training/full/Nexus-Gen.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/Nexus-Gen/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config_zero2offload.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/Nexus-Gen \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/Nexus-Gen/metadata.csv \ + --data_file_keys "image,nexus_gen_reference_image" \ + --max_pixels 262144 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "DiffSynth-Studio/Nexus-GenV2:model*.safetensors,DiffSynth-Studio/Nexus-GenV2:edit_decoder.bin,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-NexusGen-Edit_full" \ + --trainable_models "dit" \ + --extra_inputs "nexus_gen_reference_image" \ + --use_gradient_checkpointing_offload diff --git a/examples/flux/model_training/full/Step1X-Edit.sh b/examples/flux/model_training/full/Step1X-Edit.sh new file mode 100644 index 0000000000000000000000000000000000000000..b482b974e4d56ef58a520a5857fc5e631afd5848 --- /dev/null +++ b/examples/flux/model_training/full/Step1X-Edit.sh @@ -0,0 +1,16 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/Step1X-Edit/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/Step1X-Edit \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/Step1X-Edit/metadata.csv \ + --data_file_keys "image,step1x_reference_image" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "Qwen/Qwen2.5-VL-7B-Instruct:model-*.safetensors,stepfun-ai/Step1X-Edit:step1x-edit-i1258.safetensors,stepfun-ai/Step1X-Edit:vae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/Step1X-Edit_full" \ + --trainable_models "dit" \ + --extra_inputs "step1x_reference_image" \ + --use_gradient_checkpointing_offload diff --git a/examples/flux/model_training/full/accelerate_config.yaml b/examples/flux/model_training/full/accelerate_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..83280f73f315a32eccb065f351d66b4b2678759d --- /dev/null +++ b/examples/flux/model_training/full/accelerate_config.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: false + zero_stage: 2 +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 diff --git a/examples/flux/model_training/full/accelerate_config_zero2offload.yaml b/examples/flux/model_training/full/accelerate_config_zero2offload.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8a75f3d91eeae160409650b482e5383ac26b297b --- /dev/null +++ b/examples/flux/model_training/full/accelerate_config_zero2offload.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: 'cpu' + offload_param_device: 'cpu' + zero3_init_flag: false + zero_stage: 2 +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 diff --git a/examples/flux/model_training/full/accelerate_config_zero3.yaml b/examples/flux/model_training/full/accelerate_config_zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e6a8d273346c6e04e7ab97c8aa661914a0819a86 --- /dev/null +++ b/examples/flux/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 diff --git a/examples/flux/model_training/lora/FLEX.2-preview.sh b/examples/flux/model_training/lora/FLEX.2-preview.sh new file mode 100644 index 0000000000000000000000000000000000000000..ff91b3ad7bd37d002d4c0c0d44db63a1eb26db85 --- /dev/null +++ b/examples/flux/model_training/lora/FLEX.2-preview.sh @@ -0,0 +1,17 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLEX.2-preview/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLEX.2-preview \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLEX.2-preview/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "ostris/Flex.2-preview:Flex.2-preview.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLEX.2-preview_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-Fill-dev.sh b/examples/flux/model_training/lora/FLUX.1-Fill-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..92a4369d2d636650ff19612944e4f032a53710d4 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-Fill-dev.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Fill-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/metadata.csv \ + --data_file_keys "image,flux_fill_image,flux_fill_mask" \ + --max_pixels 1048576 \ + --dataset_repeat 200 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Fill-dev:flux1-fill-dev.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-Fill-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Fill-dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "flux_fill_image,flux_fill_mask" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-Kontext-dev.sh b/examples/flux/model_training/lora/FLUX.1-Kontext-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..e92818dfcfa28ae6598edc0b8cab527457fab675 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-Kontext-dev.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Kontext-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev/metadata.csv \ + --data_file_keys "image,kontext_images" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Kontext-dev:flux1-kontext-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Kontext-dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --extra_inputs "kontext_images" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-Krea-dev.sh b/examples/flux/model_training/lora/FLUX.1-Krea-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..c32ef0be2c3737734f5aa920976e21cb3d3a5466 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-Krea-dev.sh @@ -0,0 +1,17 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Krea-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Krea-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Krea-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Krea-dev:flux1-krea-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Krea-dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-Redux-dev.sh b/examples/flux/model_training/lora/FLUX.1-Redux-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..59d039c456473d6bff8c9a05aa498cc4bcd9b890 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-Redux-dev.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Redux-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev/metadata.csv \ + --data_file_keys "image,flux_redux_image" \ + --max_pixels 1048576 \ + --dataset_repeat 200 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_encoder/model.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_embedder/diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Redux-dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "flux_redux_image" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh b/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh new file mode 100644 index 0000000000000000000000000000000000000000..a09a3dccf272100a98434c612874d5c1d6f3768e --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-AttriCtrl/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-AttriCtrl \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-AttriCtrl/metadata.csv \ + --data_file_keys "image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,DiffSynth-Studio/AttriCtrl-FLUX.1-Dev:models/brightness.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-AttriCtrl_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "value_controller_inputs" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Inpainting-Beta.sh b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Inpainting-Beta.sh new file mode 100644 index 0000000000000000000000000000000000000000..ecb441d833fad9da5ad4d4b471d57a659de651fb --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Inpainting-Beta.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Inpainting-Beta/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Inpainting-Beta \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Inpainting-Beta/metadata.csv \ + --data_file_keys "image,controlnet_image,controlnet_inpaint_mask" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Inpainting-Beta_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "controlnet_image,controlnet_inpaint_mask" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Union-alpha.sh b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Union-alpha.sh new file mode 100644 index 0000000000000000000000000000000000000000..9c1a0993e23a1b9c386d517b6b71b9e6a852ebc8 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Union-alpha.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Union-alpha/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Union-alpha \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Union-alpha/metadata.csv \ + --data_file_keys "image,controlnet_image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,InstantX/FLUX.1-dev-Controlnet-Union-alpha:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Union-alpha_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "controlnet_image,controlnet_processor_id" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Upscaler.sh b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Upscaler.sh new file mode 100644 index 0000000000000000000000000000000000000000..066b21d8d994f0a256e48b40fccc9de9dd3b1ef2 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-Controlnet-Upscaler.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-Controlnet-Upscaler/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Upscaler \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-Controlnet-Upscaler/metadata.csv \ + --data_file_keys "image,controlnet_image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,jasperai/Flux.1-dev-Controlnet-Upscaler:diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-Controlnet-Upscaler_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "controlnet_image" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-EliGen.sh b/examples/flux/model_training/lora/FLUX.1-dev-EliGen.sh new file mode 100644 index 0000000000000000000000000000000000000000..e96b24fdfbd928827ff13942b3f80474287fdf32 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-EliGen.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-EliGen/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-EliGen \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-EliGen/metadata.json \ + --data_file_keys "image,eligen_entity_masks" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-EliGen_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --extra_inputs "eligen_entity_masks,eligen_entity_prompts" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-IP-Adapter.sh b/examples/flux/model_training/lora/FLUX.1-dev-IP-Adapter.sh new file mode 100644 index 0000000000000000000000000000000000000000..091a6d55c250cc7926e93445e9816b7a75d5fce8 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-IP-Adapter.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-IP-Adapter/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-IP-Adapter \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-IP-Adapter/metadata.csv \ + --data_file_keys "image,ipadapter_images" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,InstantX/FLUX.1-dev-IP-Adapter:ip-adapter.bin,google/siglip-so400m-patch14-384:model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-IP-Adapter_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "ipadapter_images" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-InfiniteYou.sh b/examples/flux/model_training/lora/FLUX.1-dev-InfiniteYou.sh new file mode 100644 index 0000000000000000000000000000000000000000..bfc2a0b92503f7b11e688c3a4be084b2c8eeb7d6 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-InfiniteYou.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev-InfiniteYou/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev-InfiniteYou \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev-InfiniteYou/metadata.csv \ + --data_file_keys "image,controlnet_image,infinityou_id_image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors,ByteDance/InfiniteYou:infu_flux_v1.0/aes_stage2/image_proj_model.bin,ByteDance/InfiniteYou:infu_flux_v1.0/aes_stage2/InfuseNetModel/*.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-InfiniteYou_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "controlnet_image,infinityou_id_image,infinityou_guidance" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev.sh b/examples/flux/model_training/lora/FLUX.1-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..3d7a73c7647465f6f0e1469dc59aa62476e19532 --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev.sh @@ -0,0 +1,17 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/Insert-Anything.sh b/examples/flux/model_training/lora/Insert-Anything.sh new file mode 100644 index 0000000000000000000000000000000000000000..9fb9cc537b4c16e30ca0686407fc1064736b51b2 --- /dev/null +++ b/examples/flux/model_training/lora/Insert-Anything.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/Insert-Anything/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/Insert-Anything \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/Insert-Anything/metadata.csv \ + --data_file_keys "image,insert_anything_source_image,insert_anything_source_mask,insert_anything_ref_image,insert_anything_ref_mask" \ + --max_pixels 1048576 \ + --dataset_repeat 200 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Fill-dev:flux1-fill-dev.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-Fill-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-Fill-dev:ae.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_encoder/model.safetensors,black-forest-labs/FLUX.1-Redux-dev:image_embedder/diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/Insert-Anything_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "insert_anything_source_image,insert_anything_source_mask,insert_anything_ref_image,insert_anything_ref_mask" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/Nexus-Gen.sh b/examples/flux/model_training/lora/Nexus-Gen.sh new file mode 100644 index 0000000000000000000000000000000000000000..ba56604e6b83679daa070399b99e838e82a5888d --- /dev/null +++ b/examples/flux/model_training/lora/Nexus-Gen.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/Nexus-Gen/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/Nexus-Gen \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/Nexus-Gen/metadata.csv \ + --data_file_keys "image,nexus_gen_reference_image" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "DiffSynth-Studio/Nexus-GenV2:model*.safetensors,DiffSynth-Studio/Nexus-GenV2:edit_decoder.bin,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-NexusGen-Edit_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --align_to_opensource_format \ + --extra_inputs "nexus_gen_reference_image" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/Step1X-Edit.sh b/examples/flux/model_training/lora/Step1X-Edit.sh new file mode 100644 index 0000000000000000000000000000000000000000..a60e165bd744a3125ff6e37cdd2145ea15e4ab08 --- /dev/null +++ b/examples/flux/model_training/lora/Step1X-Edit.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/Step1X-Edit/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/Step1X-Edit \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/Step1X-Edit/metadata.csv \ + --data_file_keys "image,step1x_reference_image" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "Qwen/Qwen2.5-VL-7B-Instruct:model-*.safetensors,stepfun-ai/Step1X-Edit:step1x-edit-i1258.safetensors,stepfun-ai/Step1X-Edit:vae.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/Step1X-Edit_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "step1x_reference_image" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/special/npu_training/FLUX.1-Kontext-dev-NPU.sh b/examples/flux/model_training/special/npu_training/FLUX.1-Kontext-dev-NPU.sh new file mode 100644 index 0000000000000000000000000000000000000000..897ecfb37b8b9c0c23576c49b99c801210d00149 --- /dev/null +++ b/examples/flux/model_training/special/npu_training/FLUX.1-Kontext-dev-NPU.sh @@ -0,0 +1,19 @@ +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export CPU_AFFINITY_CONF=1 + +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-Kontext-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config_zero2offload.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-Kontext-dev/metadata.csv \ + --data_file_keys "image,kontext_images" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-Kontext-dev:flux1-kontext-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-Kontext-dev_full" \ + --trainable_models "dit" \ + --extra_inputs "kontext_images" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/special/npu_training/FLUX.1-dev-NPU.sh b/examples/flux/model_training/special/npu_training/FLUX.1-dev-NPU.sh new file mode 100644 index 0000000000000000000000000000000000000000..fa5d5d1a36b6aee6b9e765ceb24c370931a83490 --- /dev/null +++ b/examples/flux/model_training/special/npu_training/FLUX.1-dev-NPU.sh @@ -0,0 +1,17 @@ +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export CPU_AFFINITY_CONF=1 + +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux/model_training/full/accelerate_config_zero2offload.yaml examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/special/split_training/FLUX.1-dev.sh b/examples/flux/model_training/special/split_training/FLUX.1-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..8a2c7602e4ce917f28390d3a2b074fbf6a095924 --- /dev/null +++ b/examples/flux/model_training/special/split_training/FLUX.1-dev.sh @@ -0,0 +1,38 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux/FLUX.1-dev/*" --local_dir ./data/diffsynth_example_dataset + +# Stage 1: cache deterministic preprocessing outputs. +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux/FLUX.1-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux/FLUX.1-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths 'black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors' \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/FLUX.1-dev_split_cache \ + --lora_base_model dit \ + --lora_target_modules a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp \ + --lora_rank 32 \ + --align_to_opensource_format \ + --use_gradient_checkpointing \ + --offload_models black-forest-labs/FLUX.1-dev:flux1-dev.safetensors \ + --task sft:data_process + +# Stage 2: train LoRA from the cached dataset. +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path ./models/train/FLUX.1-dev_split_cache \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths 'black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors' \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/FLUX.1-dev_split \ + --lora_base_model dit \ + --lora_target_modules a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp \ + --lora_rank 32 \ + --align_to_opensource_format \ + --use_gradient_checkpointing \ + --offload_models 'black-forest-labs/FLUX.1-dev:text_encoder_2/*.safetensors,black-forest-labs/FLUX.1-dev:ae.safetensors' \ + --task sft:train diff --git a/examples/flux/model_training/special/split_training/validate.py b/examples/flux/model_training/special/split_training/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..ecd8b9cdbf414346aaa77cbb4154ca7556ea67b9 --- /dev/null +++ b/examples/flux/model_training/special/split_training/validate.py @@ -0,0 +1,18 @@ +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_lora(pipe.dit, './models/train/FLUX.1-dev_split/epoch-4.safetensors', alpha=1) + +image = pipe(prompt="a dog", seed=0) +image.save('split_training_FLUX.1-dev.jpg') diff --git a/examples/flux/model_training/train.py b/examples/flux/model_training/train.py new file mode 100644 index 0000000000000000000000000000000000000000..33e2579500ba8d4c59dc0f19593f15b2aa333357 --- /dev/null +++ b/examples/flux/model_training/train.py @@ -0,0 +1,205 @@ +import torch, os, argparse, accelerate +from diffsynth.core import UnifiedDataset +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth.diffusion import * +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +class FluxTrainingModule(DiffusionTrainingModule): + def __init__( + self, + model_paths=None, model_id_with_origin_paths=None, + tokenizer_1_path=None, tokenizer_2_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_1_config = ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="tokenizer/") if tokenizer_1_path is None else ModelConfig(tokenizer_1_path) + tokenizer_2_config = ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="tokenizer_2/") if tokenizer_2_path is None else ModelConfig(tokenizer_2_path) + self.pipe = FluxImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_1_config=tokenizer_1_config, tokenizer_2_config=tokenizer_2_config) + self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model) + self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt) + + # Training mode + self.switch_pipe_to_training_mode( + self.pipe, trainable_models, + lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, + preset_lora_path, preset_lora_model, + task=task, + ) + + # Other configs + self.use_gradient_checkpointing = use_gradient_checkpointing + self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload + self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] + self.fp8_models = fp8_models + self.task = task + self.task_to_loss = { + "sft:data_process": lambda pipe, *args: args, + "direct_distill:data_process": lambda pipe, *args: args, + "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + "direct_distill": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi), + "direct_distill:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi), + } + + def get_pipeline_inputs(self, data): + inputs_posi = {"prompt": data["prompt"]} + inputs_nega = {"negative_prompt": ""} + inputs_shared = { + # Assume you are using this pipeline for inference, + # please fill in the input parameters. + "input_image": data["image"], + "height": data["image"].size[1], + "width": data["image"].size[0], + # Please do not modify the following parameters + # unless you clearly know what this will cause. + "cfg_scale": 1, + "embedded_guidance": 1, + "t5_sequence_length": 512, + "tiled": False, + "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 flux_parser(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser = add_general_config(parser) + parser = add_image_size_config(parser) + parser.add_argument("--tokenizer_1_path", type=str, default=None, help="Path to CLIP tokenizer.") + parser.add_argument("--tokenizer_2_path", type=str, default=None, help="Path to T5 tokenizer.") + parser.add_argument("--align_to_opensource_format", default=False, action="store_true", help="Whether to align the lora format to opensource format. Only for DiT's LoRA.") + return parser + + +def convert_lora_format(state_dict, alpha=None): + prefix_rename_dict = { + "single_blocks": "lora_unet_single_blocks", + "blocks": "lora_unet_double_blocks", + } + middle_rename_dict = { + "norm.linear": "modulation_lin", + "to_qkv_mlp": "linear1", + "proj_out": "linear2", + "norm1_a.linear": "img_mod_lin", + "norm1_b.linear": "txt_mod_lin", + "attn.a_to_qkv": "img_attn_qkv", + "attn.b_to_qkv": "txt_attn_qkv", + "attn.a_to_out": "img_attn_proj", + "attn.b_to_out": "txt_attn_proj", + "ff_a.0": "img_mlp_0", + "ff_a.2": "img_mlp_2", + "ff_b.0": "txt_mlp_0", + "ff_b.2": "txt_mlp_2", + } + suffix_rename_dict = { + "lora_B.weight": "lora_up.weight", + "lora_A.weight": "lora_down.weight", + } + state_dict_ = {} + for name, param in state_dict.items(): + names = name.split(".") + if names[-2] != "lora_A" and names[-2] != "lora_B": + names.pop(-2) + prefix = names[0] + middle = ".".join(names[2:-2]) + suffix = ".".join(names[-2:]) + block_id = names[1] + if middle not in middle_rename_dict: + continue + rename = prefix_rename_dict[prefix] + "_" + block_id + "_" + middle_rename_dict[middle] + "." + suffix_rename_dict[suffix] + state_dict_[rename] = param + if rename.endswith("lora_up.weight"): + lora_alpha = alpha if alpha is not None else param.shape[-1] + state_dict_[rename.replace("lora_up.weight", "alpha")] = torch.tensor((lora_alpha,))[0] + return state_dict_ + + +if __name__ == "__main__": + parser = flux_parser() + args = parser.parse_args() + accelerator = accelerate.Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)], + ) + dataset = UnifiedDataset( + base_path=args.dataset_base_path, + metadata_path=args.dataset_metadata_path, + repeat=args.dataset_repeat, + data_file_keys=args.data_file_keys.split(","), + main_data_operator=UnifiedDataset.default_image_operator( + base_path=args.dataset_base_path, + max_pixels=args.max_pixels, + height=args.height, + width=args.width, + height_division_factor=16, + width_division_factor=16, + ) + ) + model = FluxTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + tokenizer_1_path=args.tokenizer_1_path, + tokenizer_2_path=args.tokenizer_2_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, + state_dict_converter=convert_lora_format if args.align_to_opensource_format else lambda x:x, + enable_tensorboard_log=args.enable_tensorboard_log, + enable_swanlab_log=args.enable_swanlab_log, + swanlab_project=args.swanlab_project, + enable_wandb_log=args.enable_wandb_log, + wandb_project=args.wandb_project, + enable_csv_log=args.enable_csv_log, + ) + launcher_map = { + "sft:data_process": launch_data_process_task, + "direct_distill:data_process": launch_data_process_task, + "sft": launch_training_task, + "sft:train": launch_training_task, + "direct_distill": launch_training_task, + "direct_distill:train": launch_training_task, + } + launcher_map[args.task](accelerator, dataset, model, model_logger, args=args) diff --git a/examples/flux/model_training/validate_full/FLEX.2-preview.py b/examples/flux/model_training/validate_full/FLEX.2-preview.py new file mode 100644 index 0000000000000000000000000000000000000000..6e44fd6e568f0378c57949c37d118f90ffa46ab0 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLEX.2-preview.py @@ -0,0 +1,20 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +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"), + ], +) +state_dict = load_state_dict("models/train/FLEX.2-preview_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe(prompt="dog,white and brown dog, sitting on wall, under pink flowers", seed=0) +image.save("image_FLEX.2-preview_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-Fill-dev.py b/examples/flux/model_training/validate_full/FLUX.1-Fill-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..f9b6acfe4a8b2db4e0d00da4b94f166ca8aac479 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-Fill-dev.py @@ -0,0 +1,27 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict +from PIL import Image + + +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-Fill-dev_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe( + prompt="a white paper cup", + flux_fill_image=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/cup.png").convert("RGB"), + flux_fill_mask=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/cup_mask.png").convert("L"), + height=1632, width=1232, + seed=0, embedded_guidance=30.0, num_inference_steps=50, +) +image.save("image_FLUX.1-Fill-dev_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-Kontext-dev.py b/examples/flux/model_training/validate_full/FLUX.1-Kontext-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3604780b04aa3ee4c2b66ecb1eeca597c07d66 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-Kontext-dev.py @@ -0,0 +1,26 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-Kontext-dev_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe( + prompt="Make the dog turn its head around.", + kontext_images=Image.open("data/example_image_dataset/2.jpg").resize((768, 768)), + height=768, width=768, + seed=0 +) +image.save("image_FLUX.1-Kontext-dev_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-Krea-dev.py b/examples/flux/model_training/validate_full/FLUX.1-Krea-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..044055e271d1e4d6c59f1af17b4067bb5e5ea8b5 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-Krea-dev.py @@ -0,0 +1,20 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-Krea-dev_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe(prompt="a dog", seed=0) +image.save("image_FLUX.1-Krea-dev_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-Redux-dev.py b/examples/flux/model_training/validate_full/FLUX.1-Redux-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..2af3f4cd6c6c20f765bca033405b48c6861a4f5f --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-Redux-dev.py @@ -0,0 +1,26 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict +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="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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-Redux-dev_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe( + flux_redux_image=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev/robot.png").convert("RGB"), + embedded_guidance=2.5, num_inference_steps=50, seed=0, +) +image.save("image_FLUX.1-Redux-dev_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000000000000000000000000000000000000..74a0dcaf30bd7965670786f392d084df4acab91d --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py @@ -0,0 +1,21 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +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") + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-AttriCtrl_full/epoch-0.safetensors") +pipe.value_controller.encoders[0].load_state_dict(state_dict) + +image = pipe(prompt="a cat", seed=0, value_controller_inputs=0.1, rand_device="cuda") +image.save("image_FLUX.1-dev-AttriCtrl_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Inpainting-Beta.py b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Inpainting-Beta.py new file mode 100644 index 0000000000000000000000000000000000000000..e27d142935e13813a071810e2eb62337f1bf899e --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Inpainting-Beta.py @@ -0,0 +1,31 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +from diffsynth import load_state_dict +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-Controlnet-Inpainting-Beta_full/epoch-0.safetensors") +pipe.controlnet.models[0].load_state_dict(state_dict) + +image = pipe( + prompt="a cat sitting on a chair, wearing sunglasses", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/inpaint/image_1.jpg"), + inpaint_mask=Image.open("data/example_image_dataset/inpaint/mask.jpg"), + scale=0.9 + )], + height=1024, width=1024, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Inpainting-Beta_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Union-alpha.py b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Union-alpha.py new file mode 100644 index 0000000000000000000000000000000000000000..1db06cbca46c6d2878cc1a7e4dcca6cd01bc555e --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Union-alpha.py @@ -0,0 +1,31 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +from diffsynth import load_state_dict +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="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-Controlnet-Union-alpha_full/epoch-0.safetensors") +pipe.controlnet.models[0].load_state_dict(state_dict) + +image = pipe( + prompt="a dog", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/canny/image_1.jpg"), + scale=0.9, + processor_id="canny", + )], + height=768, width=768, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Union-alpha_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Upscaler.py b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..330e2614b8d693fa97396fad25bbaa739d150889 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-Controlnet-Upscaler.py @@ -0,0 +1,30 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +from diffsynth import load_state_dict +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="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-Controlnet-Upscaler_full/epoch-0.safetensors") +pipe.controlnet.models[0].load_state_dict(state_dict) + +image = pipe( + prompt="a dog", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/upscale/image_1.jpg"), + scale=0.9 + )], + height=768, width=768, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Upscaler_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-IP-Adapter.py b/examples/flux/model_training/validate_full/FLUX.1-dev-IP-Adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..7c15dedbe926795504c5a5b5178877ec70bbd712 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-IP-Adapter.py @@ -0,0 +1,28 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict +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="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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-IP-Adapter_full/epoch-0.safetensors") +pipe.ipadapter.load_state_dict(state_dict) + +image = pipe( + prompt="a dog", + ipadapter_images=Image.open("data/example_image_dataset/1.jpg"), + height=768, width=768, + seed=0 +) +image.save("image_FLUX.1-dev-IP-Adapter_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-InfiniteYou.py b/examples/flux/model_training/validate_full/FLUX.1-dev-InfiniteYou.py new file mode 100644 index 0000000000000000000000000000000000000000..311c5b939cfc922b59620be9d19b7e15ad087c8c --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-InfiniteYou.py @@ -0,0 +1,33 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +from diffsynth import load_state_dict +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="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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-InfiniteYou_full/epoch-0.safetensors") +state_dict_projector = {i.replace("image_proj_model.", ""): state_dict[i] for i in state_dict if i.startswith("image_proj_model.")} +pipe.image_proj_model.load_state_dict(state_dict_projector) +state_dict_controlnet = {i.replace("controlnet.models.0.", ""): state_dict[i] for i in state_dict if i.startswith("controlnet.models.0.")} +pipe.controlnet.models[0].load_state_dict(state_dict_controlnet) + +image = pipe( + prompt="a man with a red hat", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/infiniteyou/image_1.jpg"), + )], + height=1024, width=1024, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-InfiniteYou_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-LoRA-Encoder.py b/examples/flux/model_training/validate_full/FLUX.1-dev-LoRA-Encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..5c1d206ac341fcc256d7ebd081f81278321ba280 --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-LoRA-Encoder.py @@ -0,0 +1,24 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-LoRA-Encoder_full/epoch-0.safetensors") +pipe.lora_encoder.load_state_dict(state_dict) + +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. + +image = pipe(prompt="", seed=0, lora_encoder_inputs=lora) +image.save("image_FLUX.1-dev-LoRA-Encoder_full.jpg") diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev.py b/examples/flux/model_training/validate_full/FLUX.1-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..c1f9f7e4a7e87efc79badebcccb6433ec734887f --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev.py @@ -0,0 +1,20 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe(prompt="a dog", seed=0) +image.save("image_FLUX.1-dev_full.jpg") diff --git a/examples/flux/model_training/validate_full/Nexus-Gen.py b/examples/flux/model_training/validate_full/Nexus-Gen.py new file mode 100644 index 0000000000000000000000000000000000000000..1c2c2bc2bd908bf6845cfbd507291c330a39a3d2 --- /dev/null +++ b/examples/flux/model_training/validate_full/Nexus-Gen.py @@ -0,0 +1,28 @@ +import torch +from PIL import Image +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + +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"), + ], +) +state_dict = load_state_dict("models/train/FLUX.1-NexusGen-Edit_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +ref_image = Image.open("data/example_image_dataset/nexus_gen/image_1.png").convert("RGB") +prompt = "Add a pair of sunglasses." +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("NexusGen-Edit_full.jpg") diff --git a/examples/flux/model_training/validate_full/Step1X-Edit.py b/examples/flux/model_training/validate_full/Step1X-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..feaac7aabd5fd6f8b1fe211b0d5031561ebc9ffc --- /dev/null +++ b/examples/flux/model_training/validate_full/Step1X-Edit.py @@ -0,0 +1,25 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict +from PIL import Image + + +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"), + ], +) +state_dict = load_state_dict("models/train/Step1X-Edit_full/epoch-0.safetensors") +pipe.dit.load_state_dict(state_dict) + +image = pipe( + prompt="Make the dog turn its head around.", + step1x_reference_image=Image.open("data/example_image_dataset/2.jpg").resize((768, 768)), + height=768, width=768, cfg_scale=6, + seed=0 +) +image.save("image_Step1X-Edit_full.jpg") diff --git a/examples/flux/model_training/validate_lora/FLEX.2-preview.py b/examples/flux/model_training/validate_lora/FLEX.2-preview.py new file mode 100644 index 0000000000000000000000000000000000000000..a9059181d3a18f351042dc16a1052d6f4a0aac99 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLEX.2-preview.py @@ -0,0 +1,18 @@ +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="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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLEX.2-preview_lora/epoch-4.safetensors", alpha=1) + +image = pipe(prompt="dog,white and brown dog, sitting on wall, under pink flowers", seed=0) +image.save("image_FLEX.2-preview_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-Fill-dev.py b/examples/flux/model_training/validate_lora/FLUX.1-Fill-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..a1e04ae14c8ee8a0993b4b88924ce89002361944 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-Fill-dev.py @@ -0,0 +1,25 @@ +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-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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-Fill-dev_lora/epoch-0.safetensors", alpha=1) + +image = pipe( + prompt="a white paper cup", + flux_fill_image=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/cup.png").convert("RGB"), + flux_fill_mask=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Fill-dev/cup_mask.png").convert("L"), + height=1632, width=1232, + seed=0, embedded_guidance=30.0, num_inference_steps=50, +) +image.save("image_FLUX.1-Fill-dev_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-Kontext-dev.py b/examples/flux/model_training/validate_lora/FLUX.1-Kontext-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..c9e681eb551b9d83b4d49e3ab47b4cc5346c6978 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-Kontext-dev.py @@ -0,0 +1,24 @@ +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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-Kontext-dev_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="Make the dog turn its head around.", + kontext_images=Image.open("data/example_image_dataset/2.jpg").resize((768, 768)), + height=768, width=768, + seed=0 +) +image.save("image_FLUX.1-Kontext-dev_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-Krea-dev.py b/examples/flux/model_training/validate_lora/FLUX.1-Krea-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..7df61cc3eccca1972c3900f7e1c29a87526f1253 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-Krea-dev.py @@ -0,0 +1,18 @@ +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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-Krea-dev_lora/epoch-4.safetensors", alpha=1) + +image = pipe(prompt="a dog", seed=0) +image.save("image_FLUX.1-Krea-dev_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-Redux-dev.py b/examples/flux/model_training/validate_lora/FLUX.1-Redux-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..85b45f1e3469b2ddba29ceaf583a4d76545a2277 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-Redux-dev.py @@ -0,0 +1,24 @@ +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-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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-Redux-dev_lora/epoch-0.safetensors", alpha=1) + +image = pipe( + flux_redux_image=Image.open("data/diffsynth_example_dataset/flux/FLUX.1-Redux-dev/robot.png").convert("RGB"), + embedded_guidance=2.5, num_inference_steps=50, seed=0, +) +image.save("image_FLUX.1-Redux-dev_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000000000000000000000000000000000000..3fb81d25fd8244b84c514c868b04862f604e9a07 --- /dev/null +++ b/examples/flux/model_training/validate_lora/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") + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-AttriCtrl_lora/epoch-3.safetensors", alpha=1) + +image = pipe(prompt="a cat", seed=0, value_controller_inputs=0.1, rand_device="cuda") +image.save("image_FLUX.1-dev-AttriCtrl_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Inpainting-Beta.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Inpainting-Beta.py new file mode 100644 index 0000000000000000000000000000000000000000..cbedf7c1ea5b449fdef970abc4bcda8c6c5f855f --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Inpainting-Beta.py @@ -0,0 +1,29 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-Controlnet-Inpainting-Beta_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="a cat sitting on a chair, wearing sunglasses", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/inpaint/image_1.jpg"), + inpaint_mask=Image.open("data/example_image_dataset/inpaint/mask.jpg"), + scale=0.9 + )], + height=1024, width=1024, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Inpainting-Beta_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Union-alpha.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Union-alpha.py new file mode 100644 index 0000000000000000000000000000000000000000..c64c40e3ebc37f69b8de8c4f0a351bc6f9474280 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Union-alpha.py @@ -0,0 +1,29 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +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="InstantX/FLUX.1-dev-Controlnet-Union-alpha", origin_file_pattern="diffusion_pytorch_model.safetensors"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-Controlnet-Union-alpha_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="a dog", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/canny/image_1.jpg"), + scale=0.9, + processor_id="canny", + )], + height=768, width=768, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Union-alpha_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Upscaler.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..935c6fc9bdb39c838dd184ef3d29eb75dbb60998 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-Controlnet-Upscaler.py @@ -0,0 +1,28 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +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="jasperai/Flux.1-dev-Controlnet-Upscaler", origin_file_pattern="diffusion_pytorch_model.safetensors"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-Controlnet-Upscaler_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="a dog", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/upscale/image_1.jpg"), + scale=0.9 + )], + height=768, width=768, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-Controlnet-Upscaler_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-EliGen.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-EliGen.py new file mode 100644 index 0000000000000000000000000000000000000000..b252269c8e1fc8a18a3b5dfde58eeb69dee3cc84 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-EliGen.py @@ -0,0 +1,33 @@ +import torch +from PIL import Image +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_lora(pipe.dit, "models/train/FLUX.1-dev-EliGen_lora/epoch-4.safetensors", alpha=1) + +entity_prompts = ["A beautiful girl", "sign 'Entity Control'", "shorts", "shirt"] +global_prompt = "A beautiful girl wearing shirt and shorts in the street, holding a sign 'Entity Control'" +masks = [Image.open(f"data/example_image_dataset/eligen/{i}.png").convert('RGB') for i in range(len(entity_prompts))] +# generate image +image = pipe( + prompt=global_prompt, + cfg_scale=1.0, + num_inference_steps=50, + embedded_guidance=3.5, + seed=42, + height=1024, + width=1024, + eligen_entity_prompts=entity_prompts, + eligen_entity_masks=masks, +) +image.save(f"EliGen_lora.png") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-IP-Adapter.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-IP-Adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..31c295bc9fde56f20566c1f81dbcaecc2ede97f2 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-IP-Adapter.py @@ -0,0 +1,26 @@ +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-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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-IP-Adapter_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + ipadapter_images=Image.open("data/example_image_dataset/1.jpg"), + height=768, width=768, + seed=0 +) +image.save("image_FLUX.1-dev-IP-Adapter_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-InfiniteYou.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-InfiniteYou.py new file mode 100644 index 0000000000000000000000000000000000000000..9a76170b86ac2201e997b17b326231d70b65877d --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-InfiniteYou.py @@ -0,0 +1,28 @@ +import torch +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig, ControlNetInput +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="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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-InfiniteYou_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="a man with a red hat", + controlnet_inputs=[ControlNetInput( + image=Image.open("data/example_image_dataset/infiniteyou/image_1.jpg"), + )], + height=1024, width=1024, + seed=0, rand_device="cuda", +) +image.save("image_FLUX.1-dev-InfiniteYou_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev.py b/examples/flux/model_training/validate_lora/FLUX.1-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..110a65cdb82e6bca2b21224d55e2cabd50d244a3 --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev.py @@ -0,0 +1,18 @@ +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_lora(pipe.dit, "models/train/FLUX.1-dev_lora/epoch-4.safetensors", alpha=1) + +image = pipe(prompt="a dog", seed=0) +image.save("image_FLUX.1-dev_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/Insert-Anything.py b/examples/flux/model_training/validate_lora/Insert-Anything.py new file mode 100644 index 0000000000000000000000000000000000000000..1cd453536fc67ba07247829263a9e6a7a5ceae87 --- /dev/null +++ b/examples/flux/model_training/validate_lora/Insert-Anything.py @@ -0,0 +1,27 @@ +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-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, "models/train/Insert-Anything_lora/epoch-0.safetensors", alpha=1) + +image = pipe( + insert_anything_source_image=Image.open("data/diffsynth_example_dataset/flux/Insert-Anything/source_image.png").convert("RGB"), + insert_anything_source_mask=Image.open("data/diffsynth_example_dataset/flux/Insert-Anything/source_mask.png").convert("L"), + insert_anything_ref_image=Image.open("data/diffsynth_example_dataset/flux/Insert-Anything/ref_image.png").convert("RGB"), + insert_anything_ref_mask=Image.open("data/diffsynth_example_dataset/flux/Insert-Anything/ref_mask.png").convert("L"), + seed=666, embedded_guidance=30.0, num_inference_steps=50, +) +image.save("image_Insert-Anything_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/Nexus-Gen.py b/examples/flux/model_training/validate_lora/Nexus-Gen.py new file mode 100644 index 0000000000000000000000000000000000000000..447ed8f09f5857fcc3f4e4c16cfd66db4f84c768 --- /dev/null +++ b/examples/flux/model_training/validate_lora/Nexus-Gen.py @@ -0,0 +1,26 @@ +import torch +from PIL import Image +from diffsynth.pipelines.flux_image import FluxImagePipeline, ModelConfig + +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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-NexusGen-Edit_lora/epoch-4.safetensors", alpha=1) + +ref_image = Image.open("data/example_image_dataset/nexus_gen/image_1.png").convert("RGB") +prompt = "Add a pair of sunglasses." +image = pipe( + prompt=prompt, negative_prompt="", + seed=42, cfg_scale=1.0, num_inference_steps=50, + nexus_gen_reference_image=ref_image, + height=512, width=512, +) +image.save("NexusGen-Edit_lora.jpg") diff --git a/examples/flux/model_training/validate_lora/Step1X-Edit.py b/examples/flux/model_training/validate_lora/Step1X-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..e89ff9868e2795d9af79b4935f934b7c0c2e5416 --- /dev/null +++ b/examples/flux/model_training/validate_lora/Step1X-Edit.py @@ -0,0 +1,23 @@ +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="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"), + ], +) +pipe.load_lora(pipe.dit, "models/train/Step1X-Edit_lora/epoch-4.safetensors", alpha=1) + +image = pipe( + prompt="Make the dog turn its head around.", + step1x_reference_image=Image.open("data/example_image_dataset/2.jpg").resize((768, 768)), + height=768, width=768, cfg_scale=6, + seed=0 +) +image.save("image_Step1X-Edit_lora.jpg") diff --git a/examples/flux2/README.md b/examples/flux2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..459ae1458a3f21a0aad2e7933368d8177a6a2c7b --- /dev/null +++ b/examples/flux2/README.md @@ -0,0 +1,3 @@ +English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/FLUX2.html + +中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/FLUX2.html diff --git a/examples/flux2/model_inference/FLUX.2-dev.py b/examples/flux2/model_inference/FLUX.2-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..c2975367edff789f1b1e2cc1c8a5b7e439aa6d49 --- /dev/null +++ b/examples/flux2/model_inference/FLUX.2-dev.py @@ -0,0 +1,32 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from PIL import Image + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"), +) +prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom." +image = pipe(prompt, seed=42, rand_device="cuda", num_inference_steps=50) +image.save("image_FLUX.2-dev.jpg") + +prompt = "Transform the image into Japanese anime style" +edit_image = [Image.open("image_FLUX.2-dev.jpg")] +image = pipe(prompt, seed=42, rand_device="cuda", edit_image=edit_image, num_inference_steps=50, embedded_guidance=2.5) +image.save("image_FLUX.2-dev_edit.jpg") \ No newline at end of file diff --git a/examples/flux2/model_inference/FLUX.2-klein-4B.py b/examples/flux2/model_inference/FLUX.2-klein-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..2175901b96bce92d9a802d5035a63a6f2642b701 --- /dev/null +++ b/examples/flux2/model_inference/FLUX.2-klein-4B.py @@ -0,0 +1,21 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=4) +image.save("image_FLUX.2-klein-4B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=4) +image.save("image_edit_FLUX.2-klein-4B.jpg") diff --git a/examples/flux2/model_inference/FLUX.2-klein-9B.py b/examples/flux2/model_inference/FLUX.2-klein-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..b20fc2cbcf8298940f7199a9a25fa084b3989d1e --- /dev/null +++ b/examples/flux2/model_inference/FLUX.2-klein-9B.py @@ -0,0 +1,21 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=4) +image.save("image_FLUX.2-klein-9B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=4) +image.save("image_edit_FLUX.2-klein-9B.jpg") diff --git a/examples/flux2/model_inference/FLUX.2-klein-base-4B.py b/examples/flux2/model_inference/FLUX.2-klein-base-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..064e03593403d14da9329bb85a3db5c72116fd58 --- /dev/null +++ b/examples/flux2/model_inference/FLUX.2-klein-base-4B.py @@ -0,0 +1,21 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_FLUX.2-klein-base-4B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_edit_FLUX.2-klein-base-4B.jpg") diff --git a/examples/flux2/model_inference/FLUX.2-klein-base-9B.py b/examples/flux2/model_inference/FLUX.2-klein-base-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..e2e60659f49bffd69da19a02bf0f9c5cf9a5d318 --- /dev/null +++ b/examples/flux2/model_inference/FLUX.2-klein-base-9B.py @@ -0,0 +1,21 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_FLUX.2-klein-base-9B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_edit_FLUX.2-klein-base-9B.jpg") diff --git a/examples/flux2/model_inference/KleinBase4B-i2L-v2.py b/examples/flux2/model_inference/KleinBase4B-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..a8649b39a44f31ee4dafd55b134e8e3f05c149e5 --- /dev/null +++ b/examples/flux2/model_inference/KleinBase4B-i2L-v2.py @@ -0,0 +1,33 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/KleinBase4B-i2L-v2")], +) +snapshot_download("DiffSynth-Studio/KleinBase4B-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/image_1_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=42, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Aesthetic.py b/examples/flux2/model_inference/Template-KleinBase4B-Aesthetic.py new file mode 100644 index 0000000000000000000000000000000000000000..455a238bd60e70c2c2d7faa8faf6fad2e6d69af5 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Aesthetic.py @@ -0,0 +1,52 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Aesthetic")], +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 1.0, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 1.0, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_1.0.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 2.5, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 2.5, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_2.5.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Age.py b/examples/flux2/model_inference/Template-KleinBase4B-Age.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d1bef6740f470f6198c736ea81ee78203cf0a2 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Age.py @@ -0,0 +1,43 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Age")], +) +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 20}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_20.jpg") +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 50}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_50.jpg") +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 80}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_80.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Brightness.py b/examples/flux2/model_inference/Template-KleinBase4B-Brightness.py new file mode 100644 index 0000000000000000000000000000000000000000..9a25f50d89a7c6cbc6a37557936f40bb8933fd7a --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Brightness.py @@ -0,0 +1,43 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Brightness")], +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.7}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_light.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.5}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.3}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_dark.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-ContentRef.py b/examples/flux2/model_inference/Template-KleinBase4B-ContentRef.py new file mode 100644 index 0000000000000000000000000000000000000000..839e16cd6f722aa1a5bbdf3ee8ffa017ef201661 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-ContentRef.py @@ -0,0 +1,52 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image +import numpy as np + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ContentRef")], +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_1.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_2.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-ControlNet.py b/examples/flux2/model_inference/Template-KleinBase4B-ControlNet.py new file mode 100644 index 0000000000000000000000000000000000000000..d0c33a9ed69b7cc04dc9a31ffe6d1e71587581bd --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-ControlNet.py @@ -0,0 +1,54 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ControlNet")], +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone, bathed in bright sunshine.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, bathed in bright sunshine.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_sunshine.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone, surrounded by colorful magical particles.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, surrounded by colorful magical particles.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_magic.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Edit.py b/examples/flux2/model_inference/Template-KleinBase4B-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..e229f1c72a0fc5f5d60f24ef4c6e1b3220bb12e2 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Edit.py @@ -0,0 +1,54 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Edit")], +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="Put a hat on this cat.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Put a hat on this cat.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_hat.jpg") +image = template( + pipe, + prompt="Make the cat turn its head to look to the right.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Make the cat turn its head to look to the right.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_head.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Inpaint.py b/examples/flux2/model_inference/Template-KleinBase4B-Inpaint.py new file mode 100644 index 0000000000000000000000000000000000000000..c5826376be4aea69ab8a8b2f03121a2c6631925b --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Inpaint.py @@ -0,0 +1,56 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Inpaint")], +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="An orange cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + "force_inpaint": True, + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + }], +) +image.save("image_Inpaint_1.jpg") +image = template( + pipe, + prompt="A cat wearing sunglasses is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], +) +image.save("image_Inpaint_2.jpg") + diff --git a/examples/flux2/model_inference/Template-KleinBase4B-PandaMeme.py b/examples/flux2/model_inference/Template-KleinBase4B-PandaMeme.py new file mode 100644 index 0000000000000000000000000000000000000000..058816cdf8740920ea9281c765aecea891eb9a4f --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-PandaMeme.py @@ -0,0 +1,43 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-PandaMeme")], +) +image = template( + pipe, + prompt="A meme with a sleepy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_sleepy.jpg") +image = template( + pipe, + prompt="A meme with a happy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_happy.jpg") +image = template( + pipe, + prompt="A meme with a surprised expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_surprised.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Sharpness.py b/examples/flux2/model_inference/Template-KleinBase4B-Sharpness.py new file mode 100644 index 0000000000000000000000000000000000000000..e621bd7710f684ab387a01730c4fe80b2dea2ef2 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Sharpness.py @@ -0,0 +1,35 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")], +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.1}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.8}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.8.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-SoftRGB.py b/examples/flux2/model_inference/Template-KleinBase4B-SoftRGB.py new file mode 100644 index 0000000000000000000000000000000000000000..943ed6c301e4d8a476048ea2cbcdf44a2184396d --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-SoftRGB.py @@ -0,0 +1,52 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-SoftRGB")], +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 128/255, + "G": 128/255, + "B": 128/255 + }], +) +image.save("image_rgb_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 208/255, + "G": 185/255, + "B": 138/255 + }], +) +image.save("image_rgb_warm.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 94/255, + "G": 163/255, + "B": 174/255 + }], +) +image.save("image_rgb_cold.jpg") diff --git a/examples/flux2/model_inference/Template-KleinBase4B-Upscaler.py b/examples/flux2/model_inference/Template-KleinBase4B-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..d527ffb72198d12145e25d4ee954f4111a7612e8 --- /dev/null +++ b/examples/flux2/model_inference/Template-KleinBase4B-Upscaler.py @@ -0,0 +1,54 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Upscaler")], +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_1.png") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_2.png") diff --git a/examples/flux2/model_inference/TreeAdapter-KleinBase4B.py b/examples/flux2/model_inference/TreeAdapter-KleinBase4B.py new file mode 100644 index 0000000000000000000000000000000000000000..0d198661b987857f0db3f298e3a02accd4fa455b --- /dev/null +++ b/examples/flux2/model_inference/TreeAdapter-KleinBase4B.py @@ -0,0 +1,30 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! + +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/TreeAdapter-KleinBase4B", origin_file_pattern="iNaturalist/")], +) +name = "Glareola pratincola" +prompt = "A small bird with a long tail and short wings stands on sandy ground. Its plumage is light brown above, white below, with a dark collar around its neck. The background is a blurred expanse of sand." +image = template( + pipe, + seed=0, cfg_scale=4, num_inference_steps=40, + template_inputs = [{"name": name, "prompt": prompt}], + negative_template_inputs = [{"name": name}], +) +image.save("image.jpg") diff --git a/examples/flux2/model_inference_low_vram/FLUX.2-dev.py b/examples/flux2/model_inference_low_vram/FLUX.2-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..00428b14fc8f52be7d6c858c02349482a4bb9fed --- /dev/null +++ b/examples/flux2/model_inference_low_vram/FLUX.2-dev.py @@ -0,0 +1,33 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from PIL import Image + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "High resolution. A dreamy underwater portrait of a serene young woman in a flowing blue dress. Her hair floats softly around her face, strands delicately suspended in the water. Clear, shimmering light filters through, casting gentle highlights, while tiny bubbles rise around her. Her expression is calm, her features finely detailed—creating a tranquil, ethereal scene." +image = pipe(prompt, seed=42, rand_device="cuda", num_inference_steps=50) +image.save("image.jpg") + +prompt = "Transform the image into Japanese anime style" +edit_image = [Image.open("image.jpg")] +image = pipe(prompt, seed=42, rand_device="cuda", edit_image=edit_image, num_inference_steps=50, embedded_guidance=2.5) +image.save("image_edit.jpg") \ No newline at end of file diff --git a/examples/flux2/model_inference_low_vram/FLUX.2-klein-4B.py b/examples/flux2/model_inference_low_vram/FLUX.2-klein-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..b1f6f400e45c764d20ffc9af5b848180791134d5 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/FLUX.2-klein-4B.py @@ -0,0 +1,32 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=4) +image.save("image_FLUX.2-klein-4B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=4) +image.save("image_edit_FLUX.2-klein-4B.jpg") diff --git a/examples/flux2/model_inference_low_vram/FLUX.2-klein-9B.py b/examples/flux2/model_inference_low_vram/FLUX.2-klein-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..f79d8b311687075c76b1af8c6a05fc8b048e2a0c --- /dev/null +++ b/examples/flux2/model_inference_low_vram/FLUX.2-klein-9B.py @@ -0,0 +1,32 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=4) +image.save("image_FLUX.2-klein-9B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=4) +image.save("image_edit_FLUX.2-klein-9B.jpg") diff --git a/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-4B.py b/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..4538fdb32e08b9b76831faa82c52de0b4e647601 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-4B.py @@ -0,0 +1,32 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_FLUX.2-klein-base-4B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_edit_FLUX.2-klein-base-4B.jpg") diff --git a/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-9B.py b/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..65a59f682b73613cd541371bd340c69f8f8685c0 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/FLUX.2-klein-base-9B.py @@ -0,0 +1,32 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-9B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "Masterpiece, best quality. Anime-style portrait of a woman in a blue dress, underwater, surrounded by colorful bubbles." +image = pipe(prompt, seed=0, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_FLUX.2-klein-base-9B.jpg") + +prompt = "change the color of the clothes to red" +image = pipe(prompt, edit_image=[image], seed=1, rand_device="cuda", num_inference_steps=50, cfg_scale=4) +image.save("image_edit_FLUX.2-klein-base-9B.jpg") diff --git a/examples/flux2/model_inference_low_vram/KleinBase4B-i2L-v2.py b/examples/flux2/model_inference_low_vram/KleinBase4B-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..b89b4c994309e10b89724e174995c147ca6d2584 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/KleinBase4B-i2L-v2.py @@ -0,0 +1,45 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + lazy_loading=True, + model_configs=[ModelConfig(model_id="DiffSynth-Studio/KleinBase4B-i2L-v2")], +) +snapshot_download("DiffSynth-Studio/KleinBase4B-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/image_1_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=42, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Aesthetic.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Aesthetic.py new file mode 100644 index 0000000000000000000000000000000000000000..2c6f60a681364a26c680a885c843f8c52666732c --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Aesthetic.py @@ -0,0 +1,63 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Aesthetic")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 1.0, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 1.0, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_1.0.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 2.5, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": list(range(1, 180, 2)), + "lora_scales": 2.5, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_2.5.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Age.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Age.py new file mode 100644 index 0000000000000000000000000000000000000000..6fbbdf0ab5e23b4bf55a08436ca05c0fecec19df --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Age.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Age")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 20}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_20.jpg") +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 50}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_50.jpg") +image = template( + pipe, + prompt="A portrait of a woman with black hair, wearing a suit.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": 80}], + negative_template_inputs=[{"age": 45}], +) +image.save(f"image_age_80.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Brightness.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Brightness.py new file mode 100644 index 0000000000000000000000000000000000000000..8210e66be1daaf3c29849e71a0d8f4fca19f27ac --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Brightness.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Brightness")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.7}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_light.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.5}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.3}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_dark.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ContentRef.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ContentRef.py new file mode 100644 index 0000000000000000000000000000000000000000..6c0dc13c9853338415224bb6a9fd0179a5f001ce --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ContentRef.py @@ -0,0 +1,63 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image +import numpy as np + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ContentRef")], + lazy_loading=True, +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_1.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_2.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ControlNet.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ControlNet.py new file mode 100644 index 0000000000000000000000000000000000000000..3f469de9c70f8e7b832d32b8c63bf6a0b3a69aea --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-ControlNet.py @@ -0,0 +1,66 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ControlNet")], + lazy_loading=True, +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone, bathed in bright sunshine.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, bathed in bright sunshine.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_sunshine.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone, surrounded by colorful magical particles.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, surrounded by colorful magical particles.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_magic.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Edit.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..c63fb9b5aa83520f7582af792b721a87d1d2a16f --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Edit.py @@ -0,0 +1,66 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Edit")], + lazy_loading=True, +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="Put a hat on this cat.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Put a hat on this cat.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_hat.jpg") +image = template( + pipe, + prompt="Make the cat turn its head to look to the right.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Make the cat turn its head to look to the right.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_head.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Inpaint.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Inpaint.py new file mode 100644 index 0000000000000000000000000000000000000000..3106cbad2eb345d5ea1d9b3cd15761309d429c4b --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Inpaint.py @@ -0,0 +1,68 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Inpaint")], + lazy_loading=True, +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="An orange cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + "force_inpaint": True, + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + }], +) +image.save("image_Inpaint_1.jpg") +image = template( + pipe, + prompt="A cat wearing sunglasses is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], +) +image.save("image_Inpaint_2.jpg") + diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-PandaMeme.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-PandaMeme.py new file mode 100644 index 0000000000000000000000000000000000000000..3caa8e8d4baac4dd0b85aa730caacc3c8e3a4099 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-PandaMeme.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-PandaMeme")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A meme with a sleepy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_sleepy.jpg") +image = template( + pipe, + prompt="A meme with a happy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_happy.jpg") +image = template( + pipe, + prompt="A meme with a surprised expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_surprised.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Sharpness.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Sharpness.py new file mode 100644 index 0000000000000000000000000000000000000000..042f1a563ad331685d83fd7ba086509f9af1e9a4 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Sharpness.py @@ -0,0 +1,47 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.1}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.8}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.8.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-SoftRGB.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-SoftRGB.py new file mode 100644 index 0000000000000000000000000000000000000000..80c7ac80ebc5c6ffaae5b4abc6ab5ce50b05315a --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-SoftRGB.py @@ -0,0 +1,64 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-SoftRGB")], + lazy_loading=True, +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 128/255, + "G": 128/255, + "B": 128/255 + }], +) +image.save("image_rgb_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 208/255, + "G": 185/255, + "B": 138/255 + }], +) +image.save("image_rgb_warm.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 94/255, + "G": 163/255, + "B": 174/255 + }], +) +image.save("image_rgb_cold.jpg") diff --git a/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Upscaler.py b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..d303cb2f2719b31ceab91e3a127b80c79404487c --- /dev/null +++ b/examples/flux2/model_inference_low_vram/Template-KleinBase4B-Upscaler.py @@ -0,0 +1,66 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Upscaler")], + lazy_loading=True, +) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_1.png") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_2.png") diff --git a/examples/flux2/model_inference_low_vram/TreeAdapter-KleinBase4B.py b/examples/flux2/model_inference_low_vram/TreeAdapter-KleinBase4B.py new file mode 100644 index 0000000000000000000000000000000000000000..ce5f1bbd17e555d470cd1dd3204a0dd1f2bd3789 --- /dev/null +++ b/examples/flux2/model_inference_low_vram/TreeAdapter-KleinBase4B.py @@ -0,0 +1,42 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), + vram_limit=0, +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! + +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/TreeAdapter-KleinBase4B", origin_file_pattern="iNaturalist/")], + lazy_loading=True, +) +name = "Glareola pratincola" +prompt = "A small bird with a long tail and short wings stands on sandy ground. Its plumage is light brown above, white below, with a dark collar around its neck. The background is a blurred expanse of sand." +image = template( + pipe, + seed=0, cfg_scale=4, num_inference_steps=40, + template_inputs = [{"name": name, "prompt": prompt}], + negative_template_inputs = [{"name": name}], +) +image.save("image.jpg") diff --git a/examples/flux2/model_training/full/FLUX.2-klein-4B.sh b/examples/flux2/model_training/full/FLUX.2-klein-4B.sh new file mode 100644 index 0000000000000000000000000000000000000000..e201724ef9bbf97930054cb086318e099cee9deb --- /dev/null +++ b/examples/flux2/model_training/full/FLUX.2-klein-4B.sh @@ -0,0 +1,35 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-4B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-4B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-4B_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-4B_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/full/FLUX.2-klein-9B.sh b/examples/flux2/model_training/full/FLUX.2-klein-9B.sh new file mode 100644 index 0000000000000000000000000000000000000000..e34053e196cf548696f09e8dbabb75ec73a6cd3c --- /dev/null +++ b/examples/flux2/model_training/full/FLUX.2-klein-9B.sh @@ -0,0 +1,36 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-9B/*" --local_dir ./data/diffsynth_example_dataset + +# This script is tested on 8*A100 +accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-9B_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-9B_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/full/FLUX.2-klein-base-4B.sh b/examples/flux2/model_training/full/FLUX.2-klein-base-4B.sh new file mode 100644 index 0000000000000000000000000000000000000000..fe9e81f0ac7d86ab08eceda2447d13ea3ba1b62f --- /dev/null +++ b/examples/flux2/model_training/full/FLUX.2-klein-base-4B.sh @@ -0,0 +1,35 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-4B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-base-4B_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-base-4B_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh b/examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh new file mode 100644 index 0000000000000000000000000000000000000000..d7457c51c13980a64eb21d13cf3f2f0d643d03b4 --- /dev/null +++ b/examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh @@ -0,0 +1,36 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-9B/*" --local_dir ./data/diffsynth_example_dataset + +# This script is tested on 8*A100 +accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-9B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-9B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-base-9B_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-base-9B_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/full/KleinBase4B-i2L-v2.sh b/examples/flux2/model_training/full/KleinBase4B-i2L-v2.sh new file mode 100644 index 0000000000000000000000000000000000000000..ebc7be39b93dbb2237e899cba43e66c499396eb5 --- /dev/null +++ b/examples/flux2/model_training/full/KleinBase4B-i2L-v2.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/KleinBase4B-i2L-v2/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/KleinBase4B-i2L-v2 \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/KleinBase4B-i2L-v2/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/KleinBase4B-i2L-v2:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/KleinBase4B-i2L-v2_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --enable_lora_hot_loading diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Aesthetic.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Aesthetic.sh new file mode 100644 index 0000000000000000000000000000000000000000..bee97718e0a8e8fec44efbf08eb4386a24fc43c9 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Aesthetic.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Aesthetic/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Aesthetic:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Aesthetic_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --enable_lora_hot_loading diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Age.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Age.sh new file mode 100644 index 0000000000000000000000000000000000000000..d9030140897f2a5712b5e2c9511e67bca49c9e3e --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Age.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Age/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Age \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Age/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Age:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Age_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Brightness.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Brightness.sh new file mode 100644 index 0000000000000000000000000000000000000000..2506a62f8217ba7ded96d8817ae098f688ba5cbf --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Brightness.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Brightness/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Brightness:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Brightness_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-ContentRef.sh b/examples/flux2/model_training/full/Template-KleinBase4B-ContentRef.sh new file mode 100644 index 0000000000000000000000000000000000000000..52d3c2d601a74bf931fe2c2c06d7c4a080fba335 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-ContentRef.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-ContentRef/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-ContentRef \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-ContentRef/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-ContentRef:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-ContentRef_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --enable_lora_hot_loading diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-ControlNet.sh b/examples/flux2/model_training/full/Template-KleinBase4B-ControlNet.sh new file mode 100644 index 0000000000000000000000000000000000000000..ee0d23bcce09c0cbe3011199d693f1bf70f7b6c9 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-ControlNet.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-ControlNet/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-ControlNet \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-ControlNet/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-ControlNet:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-ControlNet_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Edit.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Edit.sh new file mode 100644 index 0000000000000000000000000000000000000000..536f963d05dbb41134127ce1060c3268d1c00363 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Edit.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Edit/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Edit \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Edit/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Edit:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Edit_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Inpaint.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Inpaint.sh new file mode 100644 index 0000000000000000000000000000000000000000..19ddb676dc6aece66951dc27f5a886a907dd6906 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Inpaint.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Inpaint/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Inpaint \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Inpaint/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Inpaint:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Inpaint_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-PandaMeme.sh b/examples/flux2/model_training/full/Template-KleinBase4B-PandaMeme.sh new file mode 100644 index 0000000000000000000000000000000000000000..861eb5715b4617c614b3bd83ee1d4a7097e82d08 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-PandaMeme.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-PandaMeme/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-PandaMeme \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-PandaMeme/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-PandaMeme:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-PandaMeme_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Sharpness.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Sharpness.sh new file mode 100644 index 0000000000000000000000000000000000000000..3afa3a24f8385c6be1d4b03794db949c589dd3b7 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Sharpness.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Sharpness/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Sharpness:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Sharpness_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-SoftRGB.sh b/examples/flux2/model_training/full/Template-KleinBase4B-SoftRGB.sh new file mode 100644 index 0000000000000000000000000000000000000000..add0ea8afacf5b1f51b648a392a0f859b616a49b --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-SoftRGB.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-SoftRGB/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-SoftRGB \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-SoftRGB/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-SoftRGB:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-SoftRGB_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/Template-KleinBase4B-Upscaler.sh b/examples/flux2/model_training/full/Template-KleinBase4B-Upscaler.sh new file mode 100644 index 0000000000000000000000000000000000000000..aab063bc33c255a40a3ebbe523c1dc394d76b1e0 --- /dev/null +++ b/examples/flux2/model_training/full/Template-KleinBase4B-Upscaler.sh @@ -0,0 +1,18 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Upscaler/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Upscaler \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Upscaler/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Upscaler:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Upscaler_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters diff --git a/examples/flux2/model_training/full/TreeAdapter-KleinBase4B.sh b/examples/flux2/model_training/full/TreeAdapter-KleinBase4B.sh new file mode 100644 index 0000000000000000000000000000000000000000..156fc3a8ba1f7a316f559487288326171d03ae99 --- /dev/null +++ b/examples/flux2/model_training/full/TreeAdapter-KleinBase4B.sh @@ -0,0 +1,2 @@ +# TreeAdapter-KleinBase4B is a multi-LoRA system that supports inference only. +# For training, please refer to the training scripts in Template-KleinBase4B-Aesthetic. \ No newline at end of file diff --git a/examples/flux2/model_training/full/accelerate_config.yaml b/examples/flux2/model_training/full/accelerate_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..83280f73f315a32eccb065f351d66b4b2678759d --- /dev/null +++ b/examples/flux2/model_training/full/accelerate_config.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: false + zero_stage: 2 +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 diff --git a/examples/flux2/model_training/full/accelerate_config_zero3.yaml b/examples/flux2/model_training/full/accelerate_config_zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e6a8d273346c6e04e7ab97c8aa661914a0819a86 --- /dev/null +++ b/examples/flux2/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 diff --git a/examples/flux2/model_training/lora/FLUX.2-dev.sh b/examples/flux2/model_training/lora/FLUX.2-dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..7a6c08ebbacaf516995f84be723d5c5e6a7d9622 --- /dev/null +++ b/examples/flux2/model_training/lora/FLUX.2-dev.sh @@ -0,0 +1,34 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-dev:text_encoder/*.safetensors,black-forest-labs/FLUX.2-dev:vae/diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-dev-LoRA-splited-cache" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_qkv_mlp_proj,to_out.0,to_add_out,linear_in,linear_out,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out,single_transformer_blocks.24.attn.to_out,single_transformer_blocks.25.attn.to_out,single_transformer_blocks.26.attn.to_out,single_transformer_blocks.27.attn.to_out,single_transformer_blocks.28.attn.to_out,single_transformer_blocks.29.attn.to_out,single_transformer_blocks.30.attn.to_out,single_transformer_blocks.31.attn.to_out,single_transformer_blocks.32.attn.to_out,single_transformer_blocks.33.attn.to_out,single_transformer_blocks.34.attn.to_out,single_transformer_blocks.35.attn.to_out,single_transformer_blocks.36.attn.to_out,single_transformer_blocks.37.attn.to_out,single_transformer_blocks.38.attn.to_out,single_transformer_blocks.39.attn.to_out,single_transformer_blocks.40.attn.to_out,single_transformer_blocks.41.attn.to_out,single_transformer_blocks.42.attn.to_out,single_transformer_blocks.43.attn.to_out,single_transformer_blocks.44.attn.to_out,single_transformer_blocks.45.attn.to_out,single_transformer_blocks.46.attn.to_out,single_transformer_blocks.47.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --dataset_num_workers 8 \ + --task "sft:data_process" + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path "./models/train/FLUX.2-dev-LoRA-splited-cache" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-dev:transformer/*.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-dev-LoRA-splited" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_qkv_mlp_proj,to_out.0,to_add_out,linear_in,linear_out,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out,single_transformer_blocks.24.attn.to_out,single_transformer_blocks.25.attn.to_out,single_transformer_blocks.26.attn.to_out,single_transformer_blocks.27.attn.to_out,single_transformer_blocks.28.attn.to_out,single_transformer_blocks.29.attn.to_out,single_transformer_blocks.30.attn.to_out,single_transformer_blocks.31.attn.to_out,single_transformer_blocks.32.attn.to_out,single_transformer_blocks.33.attn.to_out,single_transformer_blocks.34.attn.to_out,single_transformer_blocks.35.attn.to_out,single_transformer_blocks.36.attn.to_out,single_transformer_blocks.37.attn.to_out,single_transformer_blocks.38.attn.to_out,single_transformer_blocks.39.attn.to_out,single_transformer_blocks.40.attn.to_out,single_transformer_blocks.41.attn.to_out,single_transformer_blocks.42.attn.to_out,single_transformer_blocks.43.attn.to_out,single_transformer_blocks.44.attn.to_out,single_transformer_blocks.45.attn.to_out,single_transformer_blocks.46.attn.to_out,single_transformer_blocks.47.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --dataset_num_workers 8 \ + --task "sft:train" diff --git a/examples/flux2/model_training/lora/FLUX.2-klein-4B.sh b/examples/flux2/model_training/lora/FLUX.2-klein-4B.sh new file mode 100644 index 0000000000000000000000000000000000000000..6f4f3f4bb1f5435fc3a21963ff40f1be2b60a4c3 --- /dev/null +++ b/examples/flux2/model_training/lora/FLUX.2-klein-4B.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-4B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-4B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-4B_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-4B_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out" \ +# --lora_rank 32 \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/lora/FLUX.2-klein-9B.sh b/examples/flux2/model_training/lora/FLUX.2-klein-9B.sh new file mode 100644 index 0000000000000000000000000000000000000000..36c18027dcc53ed87eaf5a4bc8203be4bf2c7e49 --- /dev/null +++ b/examples/flux2/model_training/lora/FLUX.2-klein-9B.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-9B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-9B_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-9B_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out" \ +# --lora_rank 32 \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/lora/FLUX.2-klein-base-4B.sh b/examples/flux2/model_training/lora/FLUX.2-klein-base-4B.sh new file mode 100644 index 0000000000000000000000000000000000000000..45ef8932db0c5933895eed7d3ebf677c5a833908 --- /dev/null +++ b/examples/flux2/model_training/lora/FLUX.2-klein-base-4B.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-4B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-base-4B_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-base-4B_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out" \ +# --lora_rank 32 \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/lora/FLUX.2-klein-base-9B.sh b/examples/flux2/model_training/lora/FLUX.2-klein-base-9B.sh new file mode 100644 index 0000000000000000000000000000000000000000..aa8a32103d542b749e68caaee78a52f34dad563f --- /dev/null +++ b/examples/flux2/model_training/lora/FLUX.2-klein-base-9B.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-9B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-9B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-9B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-base-9B_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-base-9B_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out" \ +# --lora_rank 32 \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/scripts/brightness/model.py b/examples/flux2/model_training/scripts/brightness/model.py new file mode 100644 index 0000000000000000000000000000000000000000..9be2cb0a4b819ce1ab723801403e80e3ce2fc62f --- /dev/null +++ b/examples/flux2/model_training/scripts/brightness/model.py @@ -0,0 +1,62 @@ +import torch, math +from PIL import Image +import numpy as np + + +class SingleValueEncoder(torch.nn.Module): + def __init__(self, dim_in=256, dim_out=4096, length=32): + super().__init__() + self.length = length + self.prefer_value_embedder = torch.nn.Sequential(torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out)) + self.positional_embedding = torch.nn.Parameter(torch.randn(self.length, dim_out)) + + def get_timestep_embedding(self, timesteps, embedding_dim, max_period=10000): + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim + emb = timesteps[:, None].float() * torch.exp(exponent)[None, :] + emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + return emb + + def forward(self, value, dtype): + emb = self.get_timestep_embedding(value * 1000, 256).to(dtype) + emb = self.prefer_value_embedder(emb).squeeze(0) + base_embeddings = emb.expand(self.length, -1) + positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device) + learned_embeddings = base_embeddings + positional_embedding + return learned_embeddings + + +class ValueFormatModel(torch.nn.Module): + def __init__(self, num_double_blocks=5, num_single_blocks=20, dim=3072, num_heads=24, length=512): + super().__init__() + self.block_names = [f"double_{i}" for i in range(num_double_blocks)] + [f"single_{i}" for i in range(num_single_blocks)] + self.proj_k = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names}) + self.proj_v = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names}) + self.num_heads = num_heads + self.length = length + + @torch.no_grad() + def process_inputs(self, pipe, scale, **kwargs): + return {"value": torch.Tensor([scale]).to(dtype=pipe.torch_dtype, device=pipe.device)} + + def forward(self, value, **kwargs): + kv_cache = {} + for block_name in self.block_names: + k = self.proj_k[block_name](value, value.dtype) + k = k.view(1, self.length, self.num_heads, -1) + v = self.proj_v[block_name](value, value.dtype) + v = v.view(1, self.length, self.num_heads, -1) + kv_cache[block_name] = (k, v) + return {"kv_cache": kv_cache} + + +class DataAnnotator: + def __call__(self, image, **kwargs): + image = Image.open(image) + image = np.array(image) + return {"scale": image.astype(np.float32).mean() / 255} + + +TEMPLATE_MODEL = ValueFormatModel +TEMPLATE_MODEL_PATH = None # You should modify this parameter after training +TEMPLATE_DATA_PROCESSOR = DataAnnotator \ No newline at end of file diff --git a/examples/flux2/model_training/scripts/convert_base_model_to_template_model.py b/examples/flux2/model_training/scripts/convert_base_model_to_template_model.py new file mode 100644 index 0000000000000000000000000000000000000000..21fab7f62e64d1f51ed4729b809a2eca48d543b5 --- /dev/null +++ b/examples/flux2/model_training/scripts/convert_base_model_to_template_model.py @@ -0,0 +1,60 @@ +from diffsynth import load_state_dict +from safetensors.torch import save_file +import torch + + +def Flux2DiTStateDictConverter(state_dict): + rename_dict = { + "time_guidance_embed.timestep_embedder.linear_1.weight": "time_guidance_embed.timestep_embedder.0.weight", + "time_guidance_embed.timestep_embedder.linear_2.weight": "time_guidance_embed.timestep_embedder.2.weight", + "x_embedder.weight": "img_embedder.weight", + "context_embedder.weight": "txt_embedder.weight", + } + state_dict_ = {} + for name in state_dict: + if name in rename_dict: + state_dict_[rename_dict[name]] = state_dict[name] + elif name.startswith("transformer_blocks"): + if name.endswith("attn.to_q.weight"): + state_dict_[name.replace("to_q", "img_to_qkv").replace(".attn.", ".")] = torch.concat([ + state_dict[name.replace("to_q", "to_q")], + state_dict[name.replace("to_q", "to_k")], + state_dict[name.replace("to_q", "to_v")], + ], dim=0) + elif name.endswith("attn.to_k.weight") or name.endswith("attn.to_v.weight"): + continue + elif name.endswith("attn.to_out.0.weight"): + state_dict_[name.replace("attn.to_out.0.weight", "img_to_out.weight")] = state_dict[name] + elif name.endswith("attn.norm_q.weight"): + state_dict_[name.replace("attn.norm_q.weight", "img_norm_q.weight")] = state_dict[name] + elif name.endswith("attn.norm_k.weight"): + state_dict_[name.replace("attn.norm_k.weight", "img_norm_k.weight")] = state_dict[name] + elif name.endswith("attn.norm_added_q.weight"): + state_dict_[name.replace("attn.norm_added_q.weight", "txt_norm_q.weight")] = state_dict[name] + elif name.endswith("attn.norm_added_k.weight"): + state_dict_[name.replace("attn.norm_added_k.weight", "txt_norm_k.weight")] = state_dict[name] + elif name.endswith("attn.to_add_out.weight"): + state_dict_[name.replace("attn.to_add_out.weight", "txt_to_out.weight")] = state_dict[name] + elif name.endswith("attn.add_q_proj.weight"): + state_dict_[name.replace("add_q_proj", "txt_to_qkv").replace(".attn.", ".")] = torch.concat([ + state_dict[name.replace("add_q_proj", "add_q_proj")], + state_dict[name.replace("add_q_proj", "add_k_proj")], + state_dict[name.replace("add_q_proj", "add_v_proj")], + ], dim=0) + elif ".ff." in name: + state_dict_[name.replace(".ff.", ".img_ff.")] = state_dict[name] + elif ".ff_context." in name: + state_dict_[name.replace(".ff_context.", ".txt_ff.")] = state_dict[name] + elif name.endswith("attn.add_k_proj.weight") or name.endswith("attn.add_v_proj.weight"): + continue + else: + state_dict_[name] = state_dict[name] + elif name.startswith("single_transformer_blocks"): + state_dict_[name.replace(".attn.", ".")] = state_dict[name] + else: + state_dict_[name] = state_dict[name] + return state_dict_ + + +state_dict = load_state_dict("xxx.safetensors") +save_file(state_dict, "yyy.safetensors") diff --git a/examples/flux2/model_training/special/dmd2/FLUX.2-klein-base-4B-DMD2.sh b/examples/flux2/model_training/special/dmd2/FLUX.2-klein-base-4B-DMD2.sh new file mode 100644 index 0000000000000000000000000000000000000000..ea6522c67537c5dd84ae81acb2b3334c64c8dd79 --- /dev/null +++ b/examples/flux2/model_training/special/dmd2/FLUX.2-klein-base-4B-DMD2.sh @@ -0,0 +1,26 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-4B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/special/dmd2/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B/metadata.csv \ + --height 512 \ + --width 512 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-5 \ + --dmd2_fake_score_learning_rate 1e-5 \ + --dmd2_discriminator_learning_rate 1e-5 \ + --num_epochs 10 \ + --save_steps 1000 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-base-4B_dmd2" \ + --trainable_models "dit" \ + --dmd2_student_sample_steps 4 \ + --dmd2_student_sample_type sde \ + --dmd2_student_update_freq 5 \ + --dmd2_gan_loss_weight 0.03 \ + --dmd2_feature_indices 12 \ + --embedded_guidance 4 \ + --dmd2_teacher_cfg_scale 4 \ + --use_gradient_checkpointing \ No newline at end of file diff --git a/examples/flux2/model_training/special/dmd2/train.py b/examples/flux2/model_training/special/dmd2/train.py new file mode 100644 index 0000000000000000000000000000000000000000..8f62ea04a3d8828bd77d54a0a62d174321d59304 --- /dev/null +++ b/examples/flux2/model_training/special/dmd2/train.py @@ -0,0 +1,377 @@ +import torch, os, argparse, accelerate +import copy +import math + +from diffsynth.core import UnifiedDataset, gradient_checkpoint_forward +from diffsynth.diffusion import * +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig, model_fn_flux2 + +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +def _parse_int_list(value): + if value is None or value == "": + return None + return [int(i) for i in value.split(",") if i != ""] + + +def _parse_float_list(value): + if value is None or value == "": + return None + return [float(i) for i in value.split(",") if i != ""] + + +def _get_optimal_groups(num_channels): + if num_channels <= 32: + groups = max(1, num_channels // 4) + else: + groups = 32 + while groups > 1 and num_channels % groups != 0: + groups -= 1 + assert num_channels % groups == 0, f"{num_channels} not divisible by {groups}" + return groups + + +class FluxDMD2Discriminator(torch.nn.Module): + """DMD2 GAN discriminator scoring teacher hidden features at `feature_indices`. + + Only instantiated when `gan_loss_weight > 0`. + """ + + def __init__(self, feature_indices=None, num_blocks=40, inner_dim=3072): + super().__init__() + if feature_indices is None: + feature_indices = [int(num_blocks // 2)] + self.feature_indices = sorted({int(i) for i in feature_indices if 0 <= int(i) < num_blocks}) + if len(self.feature_indices) == 0: + raise ValueError("DMD2 discriminator requires at least one valid feature index.") + self.num_features = len(self.feature_indices) + self.inner_dim = inner_dim + + hidden_channels = inner_dim // 2 + self.heads = torch.nn.ModuleList([ + torch.nn.Sequential( + torch.nn.Conv2d(inner_dim, hidden_channels, kernel_size=4, stride=2, padding=1), + torch.nn.GroupNorm(_get_optimal_groups(hidden_channels), hidden_channels), + torch.nn.LeakyReLU(0.2), + torch.nn.Conv2d(hidden_channels, 1, kernel_size=1, stride=1, padding=0), + torch.nn.AdaptiveAvgPool2d((1, 1)), + torch.nn.Flatten(), + ) + for _ in self.feature_indices + ]) + + def forward(self, feats): + if not isinstance(feats, list) or len(feats) != self.num_features: + raise ValueError( + f"Expected list of {self.num_features} feature tensors, " + f"got {type(feats)} with length {len(feats) if isinstance(feats, list) else 'N/A'}." + ) + logits = [] + for head, feat in zip(self.heads, feats): + param = next(head.parameters()) + feat = feat.to(device=param.device, dtype=param.dtype) + logits.append(head(feat)) + return torch.cat(logits, dim=1) + +def model_fn_flux2_features( + dit, + latents=None, + timestep=None, + embedded_guidance=None, + prompt_embeds=None, + text_ids=None, + image_ids=None, + edit_latents=None, + edit_image_ids=None, + kv_cache=None, + extra_text_embedding=None, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + feature_indices=None, + **kwargs, +): + """Flux.2 DiT forward exposing hidden features at `feature_indices` for the discriminator.""" + feature_indices = set() if feature_indices is None else set(feature_indices) + image_seq_len = latents.shape[1] + if edit_latents is not None: + image_seq_len = latents.shape[1] + latents = torch.concat([latents, edit_latents], dim=1) + image_ids = torch.concat([image_ids, edit_image_ids], dim=1) + if embedded_guidance is None: + embedded_guidance = None + elif isinstance(embedded_guidance, torch.Tensor): + embedded_guidance = embedded_guidance.to(device=latents.device, dtype=latents.dtype).flatten() + if embedded_guidance.numel() == 1: + embedded_guidance = embedded_guidance.expand(latents.shape[0]) + elif embedded_guidance.numel() != latents.shape[0]: + raise ValueError("`embedded_guidance` must be a scalar or match the latent batch size.") + else: + embedded_guidance = torch.full((latents.shape[0],), float(embedded_guidance), device=latents.device, dtype=latents.dtype) + if extra_text_embedding is not None: + extra_text_ids = torch.zeros((1, extra_text_embedding.shape[1], 4), dtype=text_ids.dtype, device=text_ids.device) + extra_text_ids[:, :, -1] = torch.arange(prompt_embeds.shape[1], prompt_embeds.shape[1] + extra_text_embedding.shape[1]) + prompt_embeds = torch.concat([prompt_embeds, extra_text_embedding], dim=1) + text_ids = torch.concat([text_ids, extra_text_ids], dim=1) + + height, width = kwargs.get("height"), kwargs.get("width") + if height is not None and width is not None: + feature_height, feature_width = int(height) // 16, int(width) // 16 + else: + feature_height = int(math.sqrt(image_seq_len)) + feature_width = image_seq_len // feature_height if feature_height > 0 else 0 + if feature_height * feature_width != image_seq_len: + raise ValueError("Flux2 feature extraction requires height/width or square latent tokens.") + + features = [] + + def append_feature(feat): + feat = feat[:, :image_seq_len] + batch_size, _, channels = feat.shape + feat = feat.permute(0, 2, 1).reshape(batch_size, channels, feature_height, feature_width) + features.append(feat) + if len(features) == len(feature_indices): + return features + return None + + num_txt_tokens = prompt_embeds.shape[1] + timestep = timestep.to(latents.dtype) + guidance = None if embedded_guidance is None else embedded_guidance.to(latents.dtype) * 1000 + temb = dit.time_guidance_embed(timestep, guidance) + + double_stream_mod_img = dit.double_stream_modulation_img(temb) + double_stream_mod_txt = dit.double_stream_modulation_txt(temb) + single_stream_mod = dit.single_stream_modulation(temb)[0] + + hidden_states = dit.x_embedder(latents) + encoder_hidden_states = dit.context_embedder(prompt_embeds) + + if image_ids.ndim == 3: + image_ids = image_ids[0] + if text_ids.ndim == 3: + text_ids = text_ids[0] + + image_rotary_emb = dit.pos_embed(image_ids) + text_rotary_emb = dit.pos_embed(text_ids) + concat_rotary_emb = ( + torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), + torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), + ) + + for block_id, block in enumerate(dit.transformer_blocks): + encoder_hidden_states, hidden_states = gradient_checkpoint_forward( + block, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb_mod_params_img=double_stream_mod_img, + temb_mod_params_txt=double_stream_mod_txt, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=None, + kv_cache=None if kv_cache is None else kv_cache.get(f"double_{block_id}"), + ) + if block_id in feature_indices: + selected_features = append_feature(hidden_states) + if selected_features is not None: + return selected_features + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + num_double_blocks = len(dit.transformer_blocks) + + for block_id, block in enumerate(dit.single_transformer_blocks): + hidden_states = gradient_checkpoint_forward( + block, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, + hidden_states=hidden_states, + encoder_hidden_states=None, + temb_mod_params=single_stream_mod, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=None, + kv_cache=None if kv_cache is None else kv_cache.get(f"single_{block_id}"), + ) + feature_id = block_id + num_double_blocks + if feature_id in feature_indices: + selected_features = append_feature(hidden_states[:, num_txt_tokens:num_txt_tokens + image_seq_len]) + if selected_features is not None: + return selected_features + + if len(features) != len(feature_indices): + raise ValueError(f"Only collected {len(features)} feature maps for {len(feature_indices)} requested feature indices.") + return features + +def model_fn_flux2_dmd2( + pipe, + dit, + timestep, + progress_id, + num_inference_steps, + inputs_shared, + inputs_posi, + feature_indices=None, + return_features=False, +): + """Dispatcher used by `DMD2Loss`: returns flow prediction or hidden features.""" + if not return_features: + return model_fn_flux2( + dit=dit, + **inputs_shared, + **inputs_posi, + timestep=timestep, + progress_id=progress_id, + num_inference_steps=num_inference_steps, + ) + + return model_fn_flux2_features( + dit=dit, + **inputs_shared, + **inputs_posi, + timestep=timestep, + progress_id=progress_id, + num_inference_steps=num_inference_steps, + feature_indices=feature_indices, + ) + + +class Flux2DMD2TrainingModule(DiffusionTrainingModule): + def __init__(self, args, device="cpu"): + config = DMD2Config( + student_update_freq=args.dmd2_student_update_freq, + student_sample_steps=args.dmd2_student_sample_steps, + student_sample_type=args.dmd2_student_sample_type, + student_schedule=args.dmd2_student_schedule, + student_t_list=_parse_float_list(args.dmd2_student_t_list), + matching_t_min=args.dmd2_matching_t_min, + matching_t_max=args.dmd2_matching_t_max, + matching_t_sampling=args.dmd2_matching_t_sampling, + matching_t_mean=args.dmd2_matching_t_mean, + matching_t_std=args.dmd2_matching_t_std, + gan_loss_weight=args.dmd2_gan_loss_weight, + gan_r1_reg_weight=args.dmd2_gan_r1_reg_weight, + gan_r1_reg_alpha=args.dmd2_gan_r1_reg_alpha, + fake_score_learning_rate=args.dmd2_fake_score_learning_rate, + discriminator_learning_rate=args.dmd2_discriminator_learning_rate, + feature_indices=_parse_int_list(args.dmd2_feature_indices), + teacher_cfg_scale=args.dmd2_teacher_cfg_scale, + student_grad_clip_norm=args.dmd2_student_grad_clip_norm, + ) + super().__init__() + # Load models + model_configs = self.parse_model_configs(args.model_paths, args.model_id_with_origin_paths, fp8_models=args.fp8_models, offload_models=args.offload_models, device=device) + tokenizer_config = self.parse_path_or_model_id(args.tokenizer_path,default_value=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/")) + self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device,model_configs=model_configs, tokenizer_config=tokenizer_config) + + # Training mode + self.switch_pipe_to_training_mode( + self.pipe, args.trainable_models, + args.lora_base_model, args.lora_target_modules, args.lora_rank, args.lora_checkpoint, + args.preset_lora_path, args.preset_lora_model, + ) + + # DMD2 auxiliary models (frozen teacher, trainable fake_score, optional discriminator) + self.pipe.dit_teacher = copy.deepcopy(self.pipe.dit) + self.pipe.dit_fake_score = copy.deepcopy(self.pipe.dit) + self.pipe.dmd2_discriminator = None + if config.gan_loss_weight > 0: + self.pipe.dmd2_discriminator = FluxDMD2Discriminator(feature_indices=config.feature_indices) + self.pipe.dit_teacher.eval().requires_grad_(False) + self.pipe.dit_fake_score.train().requires_grad_(True) + if self.pipe.dmd2_discriminator is not None: + self.pipe.dmd2_discriminator.train().requires_grad_(True) + self.resume_from_checkpoint(args.resume_from_checkpoint, args.remove_prefix_in_ckpt) + + # Other configs + self.use_gradient_checkpointing = args.use_gradient_checkpointing + self.use_gradient_checkpointing_offload = args.use_gradient_checkpointing_offload + self.extra_inputs = args.extra_inputs.split(",") if args.extra_inputs is not None else [] + self.fp8_models = args.fp8_models + self.embedded_guidance = args.embedded_guidance + + self.dmd2_student_model_name = "dit" + self.dmd2_teacher_model_name = "dit_teacher" + self.dmd2_fake_score_model_name = "dit_fake_score" + self.dmd2_discriminator_model_name = "dmd2_discriminator" + self.dmd2_model_fn_student = model_fn_flux2_dmd2 + self.dmd2_model_fn_teacher = model_fn_flux2_dmd2 + self.dmd2_model_fn_fake_score = model_fn_flux2_dmd2 + self.dmd2_config = config + self.loss = DMD2Loss(config) + self._dmd2_student_param_names = {name for name, param in self.pipe.dit.named_parameters() if param.requires_grad} + self._dmd2_fake_score_param_names = {name for name, _ in self.pipe.dit_fake_score.named_parameters()} + + 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], + "embedded_guidance": self.embedded_guidance, + "cfg_scale": self.dmd2_config.teacher_cfg_scale, + "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, iteration=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) + return self.loss(self, inputs, 0 if iteration is None else iteration) + + def export_trainable_state_dict(self, state_dict, remove_prefix=None): + return export_dmd2_trainable_state_dict(self, state_dict, remove_prefix=remove_prefix) + +def flux2_dmd2_parser(): + parser = argparse.ArgumentParser(description="Flux.2 DMD2 training script.") + parser = add_general_config(parser) + parser = add_image_size_config(parser) + parser = add_dmd2_config(parser) + parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.") + parser.add_argument("--embedded_guidance", type=float, default=1.0, help="Flux.2 embedded guidance value.") + parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") + return parser + + +if __name__ == "__main__": + parser = flux2_dmd2_parser() + args = parser.parse_args() + args.find_unused_parameters = True + + accelerator = accelerate.Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)], + ) + dataset = UnifiedDataset( + base_path=args.dataset_base_path, + metadata_path=args.dataset_metadata_path, + repeat=args.dataset_repeat, + data_file_keys=args.data_file_keys.split(","), + main_data_operator=UnifiedDataset.default_image_operator( + base_path=args.dataset_base_path, + max_pixels=args.max_pixels, + height=args.height, + width=args.width, + height_division_factor=16, + width_division_factor=16, + ), + ) + model = Flux2DMD2TrainingModule( + args, + device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device, + ) + model_logger = ModelLogger( + args.output_path, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + enable_tensorboard_log=args.enable_tensorboard_log, + enable_swanlab_log=args.enable_swanlab_log, + swanlab_project=args.swanlab_project, + enable_wandb_log=args.enable_wandb_log, + wandb_project=args.wandb_project, + ) + launch_dmd2_training_task(accelerator, dataset, model, model_logger, args=args) diff --git a/examples/flux2/model_training/special/npu_training/FLUX.2-dev-Lora-NPU.sh b/examples/flux2/model_training/special/npu_training/FLUX.2-dev-Lora-NPU.sh new file mode 100644 index 0000000000000000000000000000000000000000..a7b4559463aaf6ffe6ca7283297084e33f6e32b5 --- /dev/null +++ b/examples/flux2/model_training/special/npu_training/FLUX.2-dev-Lora-NPU.sh @@ -0,0 +1,38 @@ +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export CPU_AFFINITY_CONF=1 + +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-dev \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-dev/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-dev:text_encoder/*.safetensors,black-forest-labs/FLUX.2-dev:vae/diffusion_pytorch_model.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-dev-LoRA-splited-cache" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_qkv_mlp_proj,to_out.0,to_add_out,linear_in,linear_out,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out,single_transformer_blocks.24.attn.to_out,single_transformer_blocks.25.attn.to_out,single_transformer_blocks.26.attn.to_out,single_transformer_blocks.27.attn.to_out,single_transformer_blocks.28.attn.to_out,single_transformer_blocks.29.attn.to_out,single_transformer_blocks.30.attn.to_out,single_transformer_blocks.31.attn.to_out,single_transformer_blocks.32.attn.to_out,single_transformer_blocks.33.attn.to_out,single_transformer_blocks.34.attn.to_out,single_transformer_blocks.35.attn.to_out,single_transformer_blocks.36.attn.to_out,single_transformer_blocks.37.attn.to_out,single_transformer_blocks.38.attn.to_out,single_transformer_blocks.39.attn.to_out,single_transformer_blocks.40.attn.to_out,single_transformer_blocks.41.attn.to_out,single_transformer_blocks.42.attn.to_out,single_transformer_blocks.43.attn.to_out,single_transformer_blocks.44.attn.to_out,single_transformer_blocks.45.attn.to_out,single_transformer_blocks.46.attn.to_out,single_transformer_blocks.47.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --dataset_num_workers 8 \ + --task "sft:data_process" + +accelerate launch --config_file examples/flux2/model_training/full/accelerate_config_zero3.yaml examples/flux2/model_training/train.py \ + --dataset_base_path "./models/train/FLUX.2-dev-LoRA-splited-cache" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-dev:transformer/*.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-dev-LoRA-splited" \ + --lora_base_model "dit" \ + --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_qkv_mlp_proj,to_out.0,to_add_out,linear_in,linear_out,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out,single_transformer_blocks.20.attn.to_out,single_transformer_blocks.21.attn.to_out,single_transformer_blocks.22.attn.to_out,single_transformer_blocks.23.attn.to_out,single_transformer_blocks.24.attn.to_out,single_transformer_blocks.25.attn.to_out,single_transformer_blocks.26.attn.to_out,single_transformer_blocks.27.attn.to_out,single_transformer_blocks.28.attn.to_out,single_transformer_blocks.29.attn.to_out,single_transformer_blocks.30.attn.to_out,single_transformer_blocks.31.attn.to_out,single_transformer_blocks.32.attn.to_out,single_transformer_blocks.33.attn.to_out,single_transformer_blocks.34.attn.to_out,single_transformer_blocks.35.attn.to_out,single_transformer_blocks.36.attn.to_out,single_transformer_blocks.37.attn.to_out,single_transformer_blocks.38.attn.to_out,single_transformer_blocks.39.attn.to_out,single_transformer_blocks.40.attn.to_out,single_transformer_blocks.41.attn.to_out,single_transformer_blocks.42.attn.to_out,single_transformer_blocks.43.attn.to_out,single_transformer_blocks.44.attn.to_out,single_transformer_blocks.45.attn.to_out,single_transformer_blocks.46.attn.to_out,single_transformer_blocks.47.attn.to_out" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --dataset_num_workers 8 \ + --initialize_model_on_cpu \ + --task "sft:train" diff --git a/examples/flux2/model_training/special/npu_training/FLUX.2-klein-9B-NPU.sh b/examples/flux2/model_training/special/npu_training/FLUX.2-klein-9B-NPU.sh new file mode 100644 index 0000000000000000000000000000000000000000..7b5c9172a7e7f2e479e783ec9c24ff54ba236aea --- /dev/null +++ b/examples/flux2/model_training/special/npu_training/FLUX.2-klein-9B-NPU.sh @@ -0,0 +1,39 @@ +# This script is tested on 8*910B(NPU) +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export CPU_AFFINITY_CONF=1 + +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-9B/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-9B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.2-klein-9B_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing + +# Edit + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch --config_file examples/flux2/model_training/full/accelerate_config.yaml examples/flux2/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-9B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-9B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-9B:vae/diffusion_pytorch_model.safetensors" \ +# --tokenizer_path "black-forest-labs/FLUX.2-klein-9B:tokenizer/" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/FLUX.2-klein-9B_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing diff --git a/examples/flux2/model_training/special/split_training/FLUX.2-klein-base-4B_lora.sh b/examples/flux2/model_training/special/split_training/FLUX.2-klein-base-4B_lora.sh new file mode 100644 index 0000000000000000000000000000000000000000..8cd53af70760e5e7b96e756abac08529df6e7ff6 --- /dev/null +++ b/examples/flux2/model_training/special/split_training/FLUX.2-klein-base-4B_lora.sh @@ -0,0 +1,38 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/FLUX.2-klein-base-4B/*" --local_dir ./data/diffsynth_example_dataset + +# Stage 1: cache deterministic preprocessing outputs. +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/FLUX.2-klein-base-4B/metadata.csv \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths 'black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors' \ + --tokenizer_path black-forest-labs/FLUX.2-klein-4B:tokenizer/ \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/FLUX.2-klein-base-4B_lora_cache \ + --lora_base_model dit \ + --lora_target_modules to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --offload_models 'black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors' \ + --task sft:data_process + +# Stage 2: train LoRA from the cached dataset. +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path ./models/train/FLUX.2-klein-base-4B_lora_cache \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths 'black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors' \ + --tokenizer_path black-forest-labs/FLUX.2-klein-4B:tokenizer/ \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/FLUX.2-klein-base-4B_lora \ + --lora_base_model dit \ + --lora_target_modules to_q,to_k,to_v,to_out.0,add_q_proj,add_k_proj,add_v_proj,to_add_out,linear_in,linear_out,to_qkv_mlp_proj,single_transformer_blocks.0.attn.to_out,single_transformer_blocks.1.attn.to_out,single_transformer_blocks.2.attn.to_out,single_transformer_blocks.3.attn.to_out,single_transformer_blocks.4.attn.to_out,single_transformer_blocks.5.attn.to_out,single_transformer_blocks.6.attn.to_out,single_transformer_blocks.7.attn.to_out,single_transformer_blocks.8.attn.to_out,single_transformer_blocks.9.attn.to_out,single_transformer_blocks.10.attn.to_out,single_transformer_blocks.11.attn.to_out,single_transformer_blocks.12.attn.to_out,single_transformer_blocks.13.attn.to_out,single_transformer_blocks.14.attn.to_out,single_transformer_blocks.15.attn.to_out,single_transformer_blocks.16.attn.to_out,single_transformer_blocks.17.attn.to_out,single_transformer_blocks.18.attn.to_out,single_transformer_blocks.19.attn.to_out \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --offload_models 'black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors' \ + --task sft:train diff --git a/examples/flux2/model_training/special/split_training/Template-KleinBase4B-Brightness.sh b/examples/flux2/model_training/special/split_training/Template-KleinBase4B-Brightness.sh new file mode 100644 index 0000000000000000000000000000000000000000..4e8154c0cc34c507101c8faa63a24a32cee994ce --- /dev/null +++ b/examples/flux2/model_training/special/split_training/Template-KleinBase4B-Brightness.sh @@ -0,0 +1,38 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Brightness/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness \ + --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness/metadata.jsonl \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --offload_models "black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Brightness:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Brightness_full_cache" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:data_process" + +accelerate launch examples/flux2/model_training/train.py \ + --dataset_base_path "./models/train/Template-KleinBase4B-Brightness_full_cache" \ + --extra_inputs "template_inputs" \ + --max_pixels 1048576 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --offload_models "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \ + --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Brightness:" \ + --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --output_path "./models/train/Template-KleinBase4B-Brightness_full" \ + --trainable_models "template_model" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:train" diff --git a/examples/flux2/model_training/special/split_training/validate.py b/examples/flux2/model_training/special/split_training/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..36115922cbd434aabd62a94ba08e43d6c6445fc4 --- /dev/null +++ b/examples/flux2/model_training/special/split_training/validate.py @@ -0,0 +1,18 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, './models/train/FLUX.2-klein-base-4B_lora/epoch-4.safetensors') +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save('split_training_FLUX.2-klein-base-4B.jpg') diff --git a/examples/flux2/model_training/train.py b/examples/flux2/model_training/train.py new file mode 100644 index 0000000000000000000000000000000000000000..d494a17d9b4a02b4be5d65e6d488eb95180ffa62 --- /dev/null +++ b/examples/flux2/model_training/train.py @@ -0,0 +1,163 @@ +import torch, os, argparse, accelerate +from diffsynth.core import UnifiedDataset +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.diffusion import * +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +class Flux2ImageTrainingModule(DiffusionTrainingModule): + def __init__( + self, + model_paths=None, model_id_with_origin_paths=None, + tokenizer_path=None, + trainable_models=None, + lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None, + preset_lora_path=None, preset_lora_model=None, + use_gradient_checkpointing=True, + use_gradient_checkpointing_offload=False, + extra_inputs=None, + fp8_models=None, + offload_models=None, + quant_options=None, + template_model_id_or_path=None, + resume_from_checkpoint=None, remove_prefix_in_ckpt=None, + enable_lora_hot_loading=False, + device="cpu", + task="sft", + ): + super().__init__() + # 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 = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/")) + self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config) + self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, use_gradient_checkpointing, use_gradient_checkpointing_offload) + self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model, remove_unnecessary_params=True) + self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt) + if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit) + + # Training mode + self.switch_pipe_to_training_mode( + self.pipe, trainable_models, + lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, + preset_lora_path, preset_lora_model, + task=task, + ) + + # Other configs + self.use_gradient_checkpointing = use_gradient_checkpointing + self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload + self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] + self.fp8_models = fp8_models + self.task = task + self.task_to_loss = { + "sft:data_process": lambda pipe, *args: args, + "direct_distill:data_process": lambda pipe, *args: args, + "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + "direct_distill": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi), + "direct_distill:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi), + } + + def get_pipeline_inputs(self, data): + inputs_posi = {"prompt": data["prompt"]} + inputs_nega = {"negative_prompt": ""} + inputs_shared = { + # Assume you are using this pipeline for inference, + # please fill in the input parameters. + "input_image": data["image"], + "height": data["image"].size[1], + "width": data["image"].size[0], + # Please do not modify the following parameters + # unless you clearly know what this will cause. + "embedded_guidance": 1.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 flux2_parser(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser = add_general_config(parser) + parser = add_image_size_config(parser) + parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.") + parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") + return parser + + +if __name__ == "__main__": + parser = flux2_parser() + args = parser.parse_args() + + accelerator = accelerate.Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)], + ) + dataset = UnifiedDataset( + base_path=args.dataset_base_path, + metadata_path=args.dataset_metadata_path, + repeat=args.dataset_repeat, + data_file_keys=args.data_file_keys.split(","), + main_data_operator=UnifiedDataset.default_image_operator( + base_path=args.dataset_base_path, + max_pixels=args.max_pixels, + height=args.height, + width=args.width, + height_division_factor=16, + width_division_factor=16, + ) + ) + model = Flux2ImageTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + tokenizer_path=args.tokenizer_path, + trainable_models=args.trainable_models, + lora_base_model=args.lora_base_model, + lora_target_modules=args.lora_target_modules, + lora_rank=args.lora_rank, + lora_checkpoint=args.lora_checkpoint, + preset_lora_path=args.preset_lora_path, + preset_lora_model=args.preset_lora_model, + use_gradient_checkpointing=args.use_gradient_checkpointing, + use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, + extra_inputs=args.extra_inputs, + fp8_models=args.fp8_models, + offload_models=args.offload_models, + quant_options=args.quant_options, + template_model_id_or_path=args.template_model_id_or_path, + resume_from_checkpoint=args.resume_from_checkpoint, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + enable_lora_hot_loading=args.enable_lora_hot_loading, + task=args.task, + device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device, + ) + model_logger = ModelLogger( + args.output_path, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + 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, + "direct_distill:data_process": launch_data_process_task, + "sft": launch_training_task, + "sft:train": launch_training_task, + "direct_distill": launch_training_task, + "direct_distill:train": launch_training_task, + } + launcher_map[args.task](accelerator, dataset, model, model_logger, args=args) \ No newline at end of file diff --git a/examples/flux2/model_training/validate_full/FLUX.2-klein-4B.py b/examples/flux2/model_training/validate_full/FLUX.2-klein-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..c5473ab8323a76a4a14d4be31c1d2e7b35e937f2 --- /dev/null +++ b/examples/flux2/model_training/validate_full/FLUX.2-klein-4B.py @@ -0,0 +1,20 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +state_dict = load_state_dict("./models/train/FLUX.2-klein-4B_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +pipe.dit.load_state_dict(state_dict) +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_full/FLUX.2-klein-9B.py b/examples/flux2/model_training/validate_full/FLUX.2-klein-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..09ac4bcd46e337c6796326b9db035c6fdb863f1e --- /dev/null +++ b/examples/flux2/model_training/validate_full/FLUX.2-klein-9B.py @@ -0,0 +1,20 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +state_dict = load_state_dict("./models/train/FLUX.2-klein-9B_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +pipe.dit.load_state_dict(state_dict) +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B-DMD2.py b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B-DMD2.py new file mode 100644 index 0000000000000000000000000000000000000000..067f019018841aee26f4c775d876d333c1b69e0f --- /dev/null +++ b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B-DMD2.py @@ -0,0 +1,20 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +state_dict = load_state_dict("./models/train/FLUX.2-klein-base-4B_dmd2/step-3000.safetensors", torch_dtype=torch.bfloat16) +pipe.dit.load_state_dict(state_dict) + +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=4, height=512, width=512) +image.save("FLUX.2-klein-base-4B_dmd2.jpg") \ No newline at end of file diff --git a/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B.py b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..95dcf9d313a51632aba268913a3b284f4f820506 --- /dev/null +++ b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-4B.py @@ -0,0 +1,20 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +state_dict = load_state_dict("./models/train/FLUX.2-klein-base-4B_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +pipe.dit.load_state_dict(state_dict) +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_full/FLUX.2-klein-base-9B.py b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..c2a192d7d1b856421dbc742e31a55295d7d75d77 --- /dev/null +++ b/examples/flux2/model_training/validate_full/FLUX.2-klein-base-9B.py @@ -0,0 +1,20 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +state_dict = load_state_dict("./models/train/FLUX.2-klein-base-9B_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +pipe.dit.load_state_dict(state_dict) +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_full/KleinBase4B-i2L-v2.py b/examples/flux2/model_training/validate_full/KleinBase4B-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..84726380bef70fc86180659871c1662ca8199c08 --- /dev/null +++ b/examples/flux2/model_training/validate_full/KleinBase4B-i2L-v2.py @@ -0,0 +1,35 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +import torch +from diffsynth import load_state_dict + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/KleinBase4B-i2L-v2")], +) +template.models[0].load_state_dict(load_state_dict("models/train/KleinBase4B-i2L-v2_full/epoch-1.safetensors")) +snapshot_download("DiffSynth-Studio/KleinBase4B-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/image_1_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=42, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Aesthetic.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Aesthetic.py new file mode 100644 index 0000000000000000000000000000000000000000..bdd66d544884dd2923c15fb2e7d8108efb8b98d8 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Aesthetic.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Aesthetic")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Aesthetic_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +image = template( + pipe, + prompt="a bird with fire", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": [1], + "lora_scales": 1.0, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": [1], + "lora_scales": 1.0, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_1.0.jpg") +image = template( + pipe, + prompt="a bird with fire", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "lora_ids": [1], + "lora_scales": 2.5, + "merge_type": "mean", + }], + negative_template_inputs = [{ + "lora_ids": [1], + "lora_scales": 2.5, + "merge_type": "mean", + }], +) +image.save("image_Aesthetic_2.5.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Age.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Age.py new file mode 100644 index 0000000000000000000000000000000000000000..d250cdefca1466d2f5da3eeb3e6b9ba511a1063a --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Age.py @@ -0,0 +1,33 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Age")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Age_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +prompt = "Half body color photograph of a single woman, head and torso with visible arms and hands resting gently in front of the body, looking directly at the camera, centered composition, colorful studio background with soft gradient of warm pastel tones, vibrant studio lighting, wearing a plain red short-sleeve t-shirt, straight black shoulder-length hair, photorealistic, high quality" +negative_age = 45 +for age in [10, 35, 70]: + print(f"Generating age {age}...") + image = template( + pipe, + prompt=prompt, + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs=[{"age": age}], + ) + image.save(f"image_age_{age}.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Brightness.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Brightness.py new file mode 100644 index 0000000000000000000000000000000000000000..7701faf75433de9accd4b7e8587e43a4282b5fbf --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Brightness.py @@ -0,0 +1,46 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Brightness")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Brightness_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.7}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_light.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.5}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.3}], + negative_template_inputs = [{"scale": 0.5}] +) +image.save("image_Brightness_dark.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-ContentRef.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-ContentRef.py new file mode 100644 index 0000000000000000000000000000000000000000..cdb9405e12cb9861b7a933c08c0ae9461349ea48 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-ContentRef.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch +from modelscope import dataset_snapshot_download +from PIL import Image +import numpy as np + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ContentRef")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-ContentRef_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_1.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_style_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.fromarray(np.zeros((1024, 1024, 3), dtype=np.uint8) + 128), + }], +) +image.save("image_ContentRef_2.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-ControlNet.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-ControlNet.py new file mode 100644 index 0000000000000000000000000000000000000000..c12b977a05d8819a72d3652cb4a20a605b065773 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-ControlNet.py @@ -0,0 +1,57 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ControlNet")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-ControlNet_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone, bathed in bright sunshine.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, bathed in bright sunshine.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_sunshine.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone, surrounded by colorful magical particles.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "A cat is sitting on a stone, surrounded by colorful magical particles.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_depth.jpg"), + "prompt": "", + }], +) +image.save("image_ControlNet_magic.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Edit.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Edit.py new file mode 100644 index 0000000000000000000000000000000000000000..5e6d2b5aeb62be43880661b8f5f44b0ae57a5345 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Edit.py @@ -0,0 +1,57 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Edit")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Edit_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="Put a hat on this cat.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Put a hat on this cat.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_hat.jpg") +image = template( + pipe, + prompt="Make the cat turn its head to look to the right.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "Make the cat turn its head to look to the right.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "prompt": "", + }], +) +image.save("image_Edit_head.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Inpaint.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Inpaint.py new file mode 100644 index 0000000000000000000000000000000000000000..5b29df72bd0e01fca6a1558bced37d17bd87bdcb --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Inpaint.py @@ -0,0 +1,59 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Inpaint")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Inpaint_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="An orange cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + "force_inpaint": True, + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_1.jpg"), + }], +) +image.save("image_Inpaint_1.jpg") +image = template( + pipe, + prompt="A cat wearing sunglasses is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_reference.jpg"), + "mask": Image.open("data/examples/templates/image_mask_2.jpg"), + }], +) +image.save("image_Inpaint_2.jpg") + diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-PandaMeme.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-PandaMeme.py new file mode 100644 index 0000000000000000000000000000000000000000..ad457b375a624eca37ccd0af4e5fcdb733dd74c3 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-PandaMeme.py @@ -0,0 +1,46 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-PandaMeme")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-PandaMeme_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +image = template( + pipe, + prompt="A meme with a sleepy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_sleepy.jpg") +image = template( + pipe, + prompt="A meme with a happy expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_happy.jpg") +image = template( + pipe, + prompt="A meme with a surprised expression.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{}], + negative_template_inputs = [{}], +) +image.save("image_PandaMeme_surprised.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Sharpness.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Sharpness.py new file mode 100644 index 0000000000000000000000000000000000000000..2a9f584cb72975db2a70a4248a6b2e0d134870ad --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Sharpness.py @@ -0,0 +1,38 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Sharpness_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.1}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.1.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"scale": 0.8}], + negative_template_inputs = [{"scale": 0.5}], +) +image.save("image_Sharpness_0.8.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-SoftRGB.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-SoftRGB.py new file mode 100644 index 0000000000000000000000000000000000000000..48865302f3169a468b6b0908f291e72abf734875 --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-SoftRGB.py @@ -0,0 +1,55 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-SoftRGB")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-SoftRGB_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 128/255, + "G": 128/255, + "B": 128/255 + }], +) +image.save("image_rgb_normal.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 208/255, + "G": 185/255, + "B": 138/255 + }], +) +image.save("image_rgb_warm.jpg") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "R": 94/255, + "G": 163/255, + "B": 174/255 + }], +) +image.save("image_rgb_cold.jpg") diff --git a/examples/flux2/model_training/validate_full/Template-KleinBase4B-Upscaler.py b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..b19c25ccfa604d0a5843fa9ce615264d8d109dca --- /dev/null +++ b/examples/flux2/model_training/validate_full/Template-KleinBase4B-Upscaler.py @@ -0,0 +1,57 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch +from modelscope import dataset_snapshot_download +from PIL import Image + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Upscaler")], +) +state_dict = load_state_dict("./models/train/Template-KleinBase4B-Upscaler_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) +dataset_snapshot_download( + "DiffSynth-Studio/examples_in_diffsynth", + allow_file_pattern=["templates/*"], + local_dir="data/examples", +) +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_512.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_1.png") +image = template( + pipe, + prompt="A cat is sitting on a stone.", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "A cat is sitting on a stone.", + }], + negative_template_inputs = [{ + "image": Image.open("data/examples/templates/image_lowres_100.jpg"), + "prompt": "", + }], +) +image.save("image_Upscaler_2.png") diff --git a/examples/flux2/model_training/validate_full/TreeAdapter-KleinBase4B.py b/examples/flux2/model_training/validate_full/TreeAdapter-KleinBase4B.py new file mode 100644 index 0000000000000000000000000000000000000000..61ffdda8bf81b1b009135d041c774af95de48ab5 --- /dev/null +++ b/examples/flux2/model_training/validate_full/TreeAdapter-KleinBase4B.py @@ -0,0 +1,34 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +from diffsynth.core import load_state_dict +import torch + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important! + +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/TreeAdapter-KleinBase4B", origin_file_pattern="iNaturalist/")], +) +state_dict = load_state_dict("./models/train/TreeAdapter-KleinBase4B_full/epoch-1.safetensors", torch_dtype=torch.bfloat16) +template.models[0].load_state_dict(state_dict) + +name = "Glareola pratincola" +prompt = "A small bird with a long tail and short wings stands on sandy ground. Its plumage is light brown above, white below, with a dark collar around its neck. The background is a blurred expanse of sand." +image = template( + pipe, + seed=0, cfg_scale=4, num_inference_steps=40, + template_inputs = [{"name": name, "prompt": prompt}], + negative_template_inputs = [{"name": name}], +) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_lora/FLUX.2-dev.py b/examples/flux2/model_training/validate_lora/FLUX.2-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..e67e2a7bfeb9b7def9f1f61f84135ed49f9a6fff --- /dev/null +++ b/examples/flux2/model_training/validate_lora/FLUX.2-dev.py @@ -0,0 +1,28 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="text_encoder/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="transformer/*.safetensors", **vram_config), + ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, "./models/train/FLUX.2-dev-LoRA-splited/epoch-4.safetensors") +prompt = "a dog" +image = pipe(prompt, seed=0) +image.save("image_FLUX.2-dev_lora.jpg") diff --git a/examples/flux2/model_training/validate_lora/FLUX.2-klein-4B.py b/examples/flux2/model_training/validate_lora/FLUX.2-klein-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..93fe2fa8535103c8f131d61b8082668a0995ad5b --- /dev/null +++ b/examples/flux2/model_training/validate_lora/FLUX.2-klein-4B.py @@ -0,0 +1,18 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, "./models/train/FLUX.2-klein-4B_lora/epoch-4.safetensors") +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_lora/FLUX.2-klein-9B.py b/examples/flux2/model_training/validate_lora/FLUX.2-klein-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..75470bc3eae6ceb45cd3284a241bb0521a46f3cb --- /dev/null +++ b/examples/flux2/model_training/validate_lora/FLUX.2-klein-9B.py @@ -0,0 +1,18 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, "./models/train/FLUX.2-klein-9B_lora/epoch-4.safetensors") +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-4B.py b/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-4B.py new file mode 100644 index 0000000000000000000000000000000000000000..6694305b2d7164b5473ebe0446da99e0b70b7d92 --- /dev/null +++ b/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-4B.py @@ -0,0 +1,18 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, "./models/train/FLUX.2-klein-base-4B_lora/epoch-4.safetensors") +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-9B.py b/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-9B.py new file mode 100644 index 0000000000000000000000000000000000000000..3551291ffc5bfde6b428fdd2b3a5d32baa17169c --- /dev/null +++ b/examples/flux2/model_training/validate_lora/FLUX.2-klein-base-9B.py @@ -0,0 +1,18 @@ +from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig +import torch + + +pipe = Flux2ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="text_encoder/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-9B", origin_file_pattern="transformer/*.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-9B", origin_file_pattern="tokenizer/"), +) +pipe.load_lora(pipe.dit, "./models/train/FLUX.2-klein-base-9B_lora/epoch-4.safetensors") +prompt = "a dog" +image = pipe(prompt=prompt, seed=0, num_inference_steps=40, cfg_scale=4, height=768, width=768) +image.save("image.jpg") diff --git a/examples/hidream_o1_image/model_inference/HiDream-O1-Image-Dev.py b/examples/hidream_o1_image/model_inference/HiDream-O1-Image-Dev.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9a2dc4d75d20a615274b017e9c1657cbb097a3 --- /dev/null +++ b/examples/hidream_o1_image/model_inference/HiDream-O1-Image-Dev.py @@ -0,0 +1,103 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline +from diffsynth.core.loader.config import ModelConfig +from diffsynth.diffusion import HiDreamO1FlashScheduler +from PIL import Image +import os +from modelscope import dataset_snapshot_download + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="./"), +) +pipe.scheduler = HiDreamO1FlashScheduler(noise_scale_start=7.5, noise_scale_end=7.5, noise_clip_std=2.5) + +# Text-to-Image +image = pipe( + prompt="medium shot, eye-level, front view. A woman is seated in an ornate bedroom, illuminated by candlelight, with a calm and composed expression. The subject is a young woman with fair skin, light brown hair styled in an updo with loose tendrils framing her face, and blue eyes. She wears a cream-colored satin robe with delicate floral embroidery and lace trim along the neckline. Her ears are adorned with pearl drop earrings. She is seated on a bed with a dark, intricately carved wooden headboard. To her left, a wooden nightstand holds three lit white candles and a candelabra with multiple lit candles in the background. The bed is covered with patterned pillows and a dark, textured blanket. The walls are paneled with dark wood and feature a large, ornate tapestry with muted earth tones. The lighting creates soft highlights on her face and robe, with warm shadows cast across the room.", + cfg_scale=1.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, +) +image.save("image.jpg") + +# Image-to-Image +image = pipe( + prompt="change her clothes to blue", + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=[image], +) +image.save("image_edit.jpg") + +dataset_snapshot_download( + dataset_id="DiffSynth-Studio/diffsynth_example_dataset", + local_dir="data/diffsynth_example_dataset", + allow_file_pattern="hidream_o1_image/HiDream-O1-Image-Dev/*", +) +# Multi-Reference Subject-Driven Personalization +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "A young boy with blonde hair stands on steps wearing light blue jeans, a white t-shirt with logo, and blue and white sneakers. He wears a brown cord necklace with beads, a black wristwatch with digital display, and carries a yellow fanny pack with white zipper. In his hand is a red boxing glove with white top, a teal plastic toy car, and a plastic toy figure of Captain America. He wears a straw hat with cream band. Natural light illuminates the scene." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, +) +image.save("ip.jpg") + +# Multi-Reference Subject-Driven Personalization with Skeleton +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP_skeleton/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "Create a realistic try-on image of the person wearing the provided clothing." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, +) +image.save("ip_skeleton.jpg") + +# Multi-Reference Subject-Driven Personalization with Layout：relative coordinates [x1, x2, y1, y2] +layout_bboxes = [[0.21, 0.44, 0.49, 0.74], [0.58, 0.80, 0.09, 0.34]] +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP_layout/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "City council members pose with relaxed smiles on a sunlit terrace, warm approachable mood, golden hour, cinematic soft glow." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, + layout_bboxes=layout_bboxes, +) +image.save("ip_layout.jpg") diff --git a/examples/hidream_o1_image/model_inference/HiDream-O1-Image.py b/examples/hidream_o1_image/model_inference/HiDream-O1-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..552f77b65d757af79f010d508350341137c1747e --- /dev/null +++ b/examples/hidream_o1_image/model_inference/HiDream-O1-Image.py @@ -0,0 +1,96 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline +from diffsynth.core.loader.config import ModelConfig +from PIL import Image +import os +from modelscope import dataset_snapshot_download + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) + +# Text-to-Image +image = pipe( + prompt="medium shot, eye-level, front view. A woman is seated in an ornate bedroom, illuminated by candlelight, with a calm and composed expression. The subject is a young woman with fair skin, light brown hair styled in an updo with loose tendrils framing her face, and blue eyes. She wears a cream-colored satin robe with delicate floral embroidery and lace trim along the neckline. Her ears are adorned with pearl drop earrings. She is seated on a bed with a dark, intricately carved wooden headboard. To her left, a wooden nightstand holds three lit white candles and a candelabra with multiple lit candles in the background. The bed is covered with patterned pillows and a dark, textured blanket. The walls are paneled with dark wood and feature a large, ornate tapestry with muted earth tones. The lighting creates soft highlights on her face and robe, with warm shadows cast across the room.", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=50, +) +image.save("image.jpg") + +# Image-to-Image +image = pipe( + prompt="change her clothes to blue", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=[image], +) +image.save("image_edit.jpg") + +dataset_snapshot_download( + dataset_id="DiffSynth-Studio/diffsynth_example_dataset", + local_dir="data/diffsynth_example_dataset", + allow_file_pattern="hidream_o1_image/HiDream-O1-Image/*", +) +# Multi-Reference Subject-Driven Personalization +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "A young boy with blonde hair stands on steps wearing light blue jeans, a white t-shirt with logo, and blue and white sneakers. He wears a brown cord necklace with beads, a black wristwatch with digital display, and carries a yellow fanny pack with white zipper. In his hand is a red boxing glove with white top, a teal plastic toy car, and a plastic toy figure of Captain America. He wears a straw hat with cream band. Natural light illuminates the scene." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, +) +image.save("ip.jpg") + +# Multi-Reference Subject-Driven Personalization with Skeleton +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP_skeleton/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "Create a realistic try-on image of the person wearing the provided clothing." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, +) +image.save("ip_skeleton.jpg") + +# Multi-Reference Subject-Driven Personalization with Layout：relative coordinates [x1, x2, y1, y2] +layout_bboxes = [[0.21, 0.44, 0.49, 0.74], [0.58, 0.80, 0.09, 0.34]] +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP_layout/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "City council members pose with relaxed smiles on a sunlit terrace, warm approachable mood, golden hour, cinematic soft glow." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, + layout_bboxes=layout_bboxes, +) +image.save("ip_layout.jpg") diff --git a/examples/hidream_o1_image/model_inference/HidreamO1-i2L-v2.py b/examples/hidream_o1_image/model_inference/HidreamO1-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..84c89bb1a8411ba640b83009aac32afbe39e3c8c --- /dev/null +++ b/examples/hidream_o1_image/model_inference/HidreamO1-i2L-v2.py @@ -0,0 +1,29 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +import torch + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors")], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/HidreamO1-i2L-v2")], +) +snapshot_download("DiffSynth-Studio/HidreamO1-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/multi_input_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image-Dev.py b/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image-Dev.py new file mode 100644 index 0000000000000000000000000000000000000000..a6974687f912d096339496556b88ae1a1927c467 --- /dev/null +++ b/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image-Dev.py @@ -0,0 +1,116 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline +from diffsynth.core.loader.config import ModelConfig +from diffsynth.diffusion import HiDreamO1FlashScheduler +from PIL import Image +import os +from modelscope import dataset_snapshot_download + + +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 = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="model-*.safetensors", **vram_config), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="./"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +pipe.scheduler = HiDreamO1FlashScheduler(noise_scale_start=7.5, noise_scale_end=7.5, noise_clip_std=2.5) + +# Text-to-Image +image = pipe( + prompt="medium shot, eye-level, front view. A woman is seated in an ornate bedroom, illuminated by candlelight, with a calm and composed expression. The subject is a young woman with fair skin, light brown hair styled in an updo with loose tendrils framing her face, and blue eyes. She wears a cream-colored satin robe with delicate floral embroidery and lace trim along the neckline. Her ears are adorned with pearl drop earrings. She is seated on a bed with a dark, intricately carved wooden headboard. To her left, a wooden nightstand holds three lit white candles and a candelabra with multiple lit candles in the background. The bed is covered with patterned pillows and a dark, textured blanket. The walls are paneled with dark wood and feature a large, ornate tapestry with muted earth tones. The lighting creates soft highlights on her face and robe, with warm shadows cast across the room.", + cfg_scale=1.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, +) +image.save("image.jpg") + +# Image-to-Image +image = pipe( + prompt="change her clothes to blue", + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=[image], +) +image.save("image_edit.jpg") + +dataset_snapshot_download( + dataset_id="DiffSynth-Studio/diffsynth_example_dataset", + local_dir="data/diffsynth_example_dataset", + allow_file_pattern="hidream_o1_image/HiDream-O1-Image-Dev/*", +) +# Multi-Reference Subject-Driven Personalization +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "A young boy with blonde hair stands on steps wearing light blue jeans, a white t-shirt with logo, and blue and white sneakers. He wears a brown cord necklace with beads, a black wristwatch with digital display, and carries a yellow fanny pack with white zipper. In his hand is a red boxing glove with white top, a teal plastic toy car, and a plastic toy figure of Captain America. He wears a straw hat with cream band. Natural light illuminates the scene." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, +) +image.save("ip.jpg") + +# Multi-Reference Subject-Driven Personalization with Skeleton +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP_skeleton/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "Create a realistic try-on image of the person wearing the provided clothing." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, +) +image.save("ip_skeleton.jpg") + +# Multi-Reference Subject-Driven Personalization with Layout：relative coordinates [x1, x2, y1, y2] +layout_bboxes = [[0.21, 0.44, 0.49, 0.74], [0.58, 0.80, 0.09, 0.34]] +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/IP_layout/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "City council members pose with relaxed smiles on a sunlit terrace, warm approachable mood, golden hour, cinematic soft glow." +image = pipe( + prompt=prompt, + cfg_scale=1.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, + edit_image=ref_images, + layout_bboxes=layout_bboxes, +) +image.save("ip_layout.jpg") diff --git a/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image.py b/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..2b23c7691fb8f2661f983efc3e865f41e51bbfb2 --- /dev/null +++ b/examples/hidream_o1_image/model_inference_low_vram/HiDream-O1-Image.py @@ -0,0 +1,109 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline +from diffsynth.core.loader.config import ModelConfig +from modelscope import dataset_snapshot_download +from PIL import Image +import os + + +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 = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors", **vram_config), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +# Text-to-Image +image = pipe( + prompt="medium shot, eye-level, front view. A woman is seated in an ornate bedroom, illuminated by candlelight, with a calm and composed expression. The subject is a young woman with fair skin, light brown hair styled in an updo with loose tendrils framing her face, and blue eyes. She wears a cream-colored satin robe with delicate floral embroidery and lace trim along the neckline. Her ears are adorned with pearl drop earrings. She is seated on a bed with a dark, intricately carved wooden headboard. To her left, a wooden nightstand holds three lit white candles and a candelabra with multiple lit candles in the background. The bed is covered with patterned pillows and a dark, textured blanket. The walls are paneled with dark wood and feature a large, ornate tapestry with muted earth tones. The lighting creates soft highlights on her face and robe, with warm shadows cast across the room.", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=50, +) +image.save("image.jpg") + +# Image-to-Image +image = pipe( + prompt="change her clothes to blue", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=[image], +) +image.save("image_edit.jpg") + +dataset_snapshot_download( + dataset_id="DiffSynth-Studio/diffsynth_example_dataset", + local_dir="data/diffsynth_example_dataset", + allow_file_pattern="hidream_o1_image/HiDream-O1-Image/*", +) +# Multi-Reference Subject-Driven Personalization +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "A young boy with blonde hair stands on steps wearing light blue jeans, a white t-shirt with logo, and blue and white sneakers. He wears a brown cord necklace with beads, a black wristwatch with digital display, and carries a yellow fanny pack with white zipper. In his hand is a red boxing glove with white top, a teal plastic toy car, and a plastic toy figure of Captain America. He wears a straw hat with cream band. Natural light illuminates the scene." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, +) +image.save("ip.jpg") + +# Multi-Reference Subject-Driven Personalization with Skeleton +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP_skeleton/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "Create a realistic try-on image of the person wearing the provided clothing." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, +) +image.save("ip_skeleton.jpg") + +# Multi-Reference Subject-Driven Personalization with Layout：relative coordinates [x1, x2, y1, y2] +layout_bboxes = [[0.21, 0.44, 0.49, 0.74], [0.58, 0.80, 0.09, 0.34]] +ref_image_dir = "./data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/IP_layout/" +ref_images = [Image.open(os.path.join(ref_image_dir, f)).convert("RGB") for f in os.listdir(ref_image_dir) if f.endswith(".jpg")] +prompt = "City council members pose with relaxed smiles on a sunlit terrace, warm approachable mood, golden hour, cinematic soft glow." +image = pipe( + prompt=prompt, + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=43, + num_inference_steps=50, + edit_image=ref_images, + layout_bboxes=layout_bboxes, +) +image.save("ip_layout.jpg") diff --git a/examples/hidream_o1_image/model_inference_low_vram/HidreamO1-i2L-v2.py b/examples/hidream_o1_image/model_inference_low_vram/HidreamO1-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..794ad4633cad0307f863582101805881c72363e2 --- /dev/null +++ b/examples/hidream_o1_image/model_inference_low_vram/HidreamO1-i2L-v2.py @@ -0,0 +1,41 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +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 = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors", **vram_config),], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + lazy_loading=True, + model_configs=[ModelConfig(model_id="DiffSynth-Studio/HidreamO1-i2L-v2")], +) +snapshot_download("DiffSynth-Studio/HidreamO1-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/multi_input_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/hidream_o1_image/model_training/full/HiDream-O1-Image-Dev.sh b/examples/hidream_o1_image/model_training/full/HiDream-O1-Image-Dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..2883b6bc514cdb3e2482b49184c6c0aa8a57bafd --- /dev/null +++ b/examples/hidream_o1_image/model_training/full/HiDream-O1-Image-Dev.sh @@ -0,0 +1,35 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HiDream-O1-Image-Dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/hidream_o1_image/model_training/full/accelerate_config_zero3.yaml examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/metadata.csv \ + --max_pixels 4194304 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image-Dev:model-*.safetensors" \ + --processor_config "HiDream-ai/HiDream-O1-Image-Dev:./" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/HiDream-O1-Image-Dev_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --noise_scale 7.5 + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/hidream_o1_image/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 4194304 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image-Dev:model-*.safetensors" \ +# --processor_config "HiDream-ai/HiDream-O1-Image-Dev:./" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/HiDream-O1-Image-Dev_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing \ +# --noise_scale 7.5 diff --git a/examples/hidream_o1_image/model_training/full/HiDream-O1-Image.sh b/examples/hidream_o1_image/model_training/full/HiDream-O1-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..e14047c0a268b5f298d27f3b5d9c98355a4b8a13 --- /dev/null +++ b/examples/hidream_o1_image/model_training/full/HiDream-O1-Image.sh @@ -0,0 +1,35 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HiDream-O1-Image/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/hidream_o1_image/model_training/full/accelerate_config_zero3.yaml examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/metadata.csv \ + --max_pixels 4194304 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image:model-*.safetensors" \ + --processor_config "HiDream-ai/HiDream-O1-Image:./" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/HiDream-O1-Image_full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --noise_scale 8.0 + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/hidream_o1_image/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 4194304 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image:model-*.safetensors" \ +# --processor_config "HiDream-ai/HiDream-O1-Image:./" \ +# --learning_rate 1e-5 \ +# --num_epochs 2 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/HiDream-O1-Image_full" \ +# --trainable_models "dit" \ +# --use_gradient_checkpointing \ +# --noise_scale 8.0 diff --git a/examples/hidream_o1_image/model_training/full/HidreamO1-i2L-v2.sh b/examples/hidream_o1_image/model_training/full/HidreamO1-i2L-v2.sh new file mode 100644 index 0000000000000000000000000000000000000000..eab0736529e4c5f79937002ad262579f3fd6b4ce --- /dev/null +++ b/examples/hidream_o1_image/model_training/full/HidreamO1-i2L-v2.sh @@ -0,0 +1,19 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HidreamO1-i2L-v2/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HidreamO1-i2L-v2 \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HidreamO1-i2L-v2/metadata.jsonl \ + --max_pixels 4194304 \ + --dataset_repeat 400 \ + --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image:model-*.safetensors" \ + --processor_config "HiDream-ai/HiDream-O1-Image:./" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --output_path "./models/train/HidreamO1-i2L-v2_full" \ + --use_gradient_checkpointing \ + --noise_scale 8.0 \ + --enable_lora_hot_loading \ + --template_model_id_or_path "DiffSynth-Studio/HidreamO1-i2L-v2:" \ + --extra_inputs "template_inputs" \ + --remove_prefix_in_ckpt "pipe.template_model." \ + --trainable_models "template_model" # Use `template_model.emb2lora` to freeze the image encoder. diff --git a/examples/hidream_o1_image/model_training/full/accelerate_config_zero3.yaml b/examples/hidream_o1_image/model_training/full/accelerate_config_zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e6a8d273346c6e04e7ab97c8aa661914a0819a86 --- /dev/null +++ b/examples/hidream_o1_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 diff --git a/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image-Dev.sh b/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image-Dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..a3c3f1a14a2a529dbcc95b91a9971cdcb7627c98 --- /dev/null +++ b/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image-Dev.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HiDream-O1-Image-Dev/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image-Dev/metadata.csv \ + --max_pixels 4194304 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image-Dev:model-*.safetensors" \ + --processor_config "HiDream-ai/HiDream-O1-Image-Dev:./" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --lora_rank 32 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/HiDream-O1-Image-Dev_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2" \ + --use_gradient_checkpointing \ + --noise_scale 7.5 + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/hidream_o1_image/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 4194304 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image-Dev:model-*.safetensors" \ +# --processor_config "HiDream-ai/HiDream-O1-Image-Dev:./" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --lora_rank 32 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/HiDream-O1-Image-Dev_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2" \ +# --use_gradient_checkpointing \ +# --noise_scale 7.5 diff --git a/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image.sh b/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..81e9ad4cdb7ea4bdf04712caf346167ac7e775b9 --- /dev/null +++ b/examples/hidream_o1_image/model_training/lora/HiDream-O1-Image.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HiDream-O1-Image/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/metadata.csv \ + --max_pixels 4194304 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image:model-*.safetensors" \ + --processor_config "HiDream-ai/HiDream-O1-Image:./" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --lora_rank 32 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/HiDream-O1-Image_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2" \ + --use_gradient_checkpointing \ + --noise_scale 8.0 + +# modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "qwen_image/Qwen-Image-Edit-2511/*" --local_dir ./data/diffsynth_example_dataset + +# accelerate launch examples/hidream_o1_image/model_training/train.py \ +# --dataset_base_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511 \ +# --dataset_metadata_path data/diffsynth_example_dataset/qwen_image/Qwen-Image-Edit-2511/metadata.json \ +# --data_file_keys "image,edit_image" \ +# --extra_inputs "edit_image" \ +# --max_pixels 4194304 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "HiDream-ai/HiDream-O1-Image:model-*.safetensors" \ +# --processor_config "HiDream-ai/HiDream-O1-Image:./" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --lora_rank 32 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/HiDream-O1-Image_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2" \ +# --use_gradient_checkpointing \ +# --noise_scale 8.0 diff --git a/examples/hidream_o1_image/model_training/special/split_training/HiDream-O1-Image.sh b/examples/hidream_o1_image/model_training/special/split_training/HiDream-O1-Image.sh new file mode 100644 index 0000000000000000000000000000000000000000..07059172d3125258263386aa442d237257f187b5 --- /dev/null +++ b/examples/hidream_o1_image/model_training/special/split_training/HiDream-O1-Image.sh @@ -0,0 +1,39 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "hidream_o1_image/HiDream-O1-Image/*" --local_dir ./data/diffsynth_example_dataset + +# Stage 1: cache deterministic preprocessing outputs. +accelerate launch examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image \ + --dataset_metadata_path data/diffsynth_example_dataset/hidream_o1_image/HiDream-O1-Image/metadata.csv \ + --max_pixels 4194304 \ + --dataset_repeat 1 \ + --model_id_with_origin_paths 'HiDream-ai/HiDream-O1-Image:model-*.safetensors' \ + --processor_config HiDream-ai/HiDream-O1-Image:./ \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --lora_rank 32 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/HiDream-O1-Image_split_cache \ + --lora_base_model dit \ + --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2 \ + --use_gradient_checkpointing \ + --noise_scale 8.0 \ + --offload_models 'HiDream-ai/HiDream-O1-Image:model-*.safetensors' \ + --task sft:data_process + +# Stage 2: train LoRA from the cached dataset. +accelerate launch examples/hidream_o1_image/model_training/train.py \ + --dataset_base_path ./models/train/HiDream-O1-Image_split_cache \ + --max_pixels 4194304 \ + --dataset_repeat 50 \ + --model_id_with_origin_paths 'HiDream-ai/HiDream-O1-Image:model-*.safetensors' \ + --processor_config HiDream-ai/HiDream-O1-Image:./ \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --lora_rank 32 \ + --remove_prefix_in_ckpt pipe.dit. \ + --output_path ./models/train/HiDream-O1-Image_split \ + --lora_base_model dit \ + --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,attn.qkv,attn.proj,mlp.linear_fc1,mlp.linear_fc2 \ + --use_gradient_checkpointing \ + --noise_scale 8.0 \ + --task sft:train diff --git a/examples/hidream_o1_image/model_training/special/split_training/validate.py b/examples/hidream_o1_image/model_training/special/split_training/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..d3c2d36f87a74f9e92ff59a9a6f5e6b0325bb180 --- /dev/null +++ b/examples/hidream_o1_image/model_training/special/split_training/validate.py @@ -0,0 +1,23 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) +pipe.load_lora(pipe.dit, './models/train/HiDream-O1-Image_split/epoch-4.safetensors') +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=50, +) +image.save('split_training_HiDream-O1-Image.jpg') diff --git a/examples/hidream_o1_image/model_training/train.py b/examples/hidream_o1_image/model_training/train.py new file mode 100644 index 0000000000000000000000000000000000000000..561a5cd4168e018538aafd67e27a1dc7d3628a84 --- /dev/null +++ b/examples/hidream_o1_image/model_training/train.py @@ -0,0 +1,158 @@ +# HiDream-O1-Image training module for DiffSynth-Studio. + +import torch, os, argparse, accelerate +import numpy as np +from diffsynth.core import UnifiedDataset +from diffsynth.core.data.operators import * +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig +from diffsynth.diffusion import * +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +class HiDreamO1ImageTrainingModule(DiffusionTrainingModule): + def __init__( + self, + model_paths=None, model_id_with_origin_paths=None, + processor_config=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", + noise_scale=8.0, + template_model_id_or_path=None, + enable_lora_hot_loading=False, + ): + super().__init__() + model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, quant_options=quant_options, device=device) + processor_config = self.parse_path_or_model_id(processor_config, default_value=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="./")) + self.pipe = HiDreamO1ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, processor_config=processor_config) + self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, use_gradient_checkpointing, use_gradient_checkpointing_offload) + self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model) + self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt) + if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit) + self.switch_pipe_to_training_mode( + self.pipe, trainable_models, + lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, + preset_lora_path, preset_lora_model, + task=task, + ) + + self.use_gradient_checkpointing = use_gradient_checkpointing + self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload + self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] + self.fp8_models = fp8_models + self.task = task + self.noise_scale = noise_scale + self.task_to_loss = { + "sft:data_process": lambda pipe, *args: args, + "sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + "sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi), + } + + def get_pipeline_inputs(self, data): + inputs_posi = {"prompt": data["prompt"]} + inputs_nega = {"negative_prompt": " "} + image = data["image"] + inputs_shared = { + "input_image": image, + "height": image.size[1], + "width": image.size[0], + "cfg_scale": 1, + "rand_device": self.pipe.device, + "noise_scale": self.noise_scale, + "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 hidream_o1_image_parser(): + parser = argparse.ArgumentParser(description="HiDream-O1-Image training.") + parser = add_general_config(parser) + parser = add_image_size_config(parser) + parser.add_argument("--processor_config", type=str, default=None, help="Path to processor config.") + parser.add_argument("--noise_scale", type=float, default=8.0, help="Noise scale factor.") + parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") + return parser + + +if __name__ == "__main__": + parser = hidream_o1_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)], + ) + from diffsynth.models.hidream_common import PATCH_SIZE + dataset = UnifiedDataset( + base_path=args.dataset_base_path, + metadata_path=args.dataset_metadata_path, + repeat=args.dataset_repeat, + data_file_keys=args.data_file_keys.split(","), + main_data_operator=UnifiedDataset.default_image_operator( + base_path=args.dataset_base_path, + max_pixels=args.max_pixels, + height=args.height, + width=args.width, + height_division_factor=PATCH_SIZE, + width_division_factor=PATCH_SIZE, + ), + ) + model = HiDreamO1ImageTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + processor_config=args.processor_config, + trainable_models=args.trainable_models, + lora_base_model=args.lora_base_model, + lora_target_modules=args.lora_target_modules, + lora_rank=args.lora_rank, + lora_checkpoint=args.lora_checkpoint, + preset_lora_path=args.preset_lora_path, + preset_lora_model=args.preset_lora_model, + use_gradient_checkpointing=args.use_gradient_checkpointing, + use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, + extra_inputs=args.extra_inputs, + fp8_models=args.fp8_models, + offload_models=args.offload_models, + quant_options=args.quant_options, + resume_from_checkpoint=args.resume_from_checkpoint, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + task=args.task, + device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device, + noise_scale=args.noise_scale, + template_model_id_or_path=args.template_model_id_or_path, + enable_lora_hot_loading=args.enable_lora_hot_loading, + ) + 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) diff --git a/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image-Dev.py b/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image-Dev.py new file mode 100644 index 0000000000000000000000000000000000000000..2410f36ffda70777cc379e0ce6dc2e86a9da6443 --- /dev/null +++ b/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image-Dev.py @@ -0,0 +1,27 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="./"), +) +pipe.scheduler = HiDreamO1FlashScheduler(noise_scale_start=7.5, noise_scale_end=7.5, noise_clip_std=2.5) + +state_dict = torch.load("./models/train/HiDream-O1-Image-Dev_full/epoch-1.safetensors", map_location="cuda", weights_only=False) +pipe.dit.load_state_dict(state_dict, strict=False) +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + cfg_scale=1.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, +) +image.save("image.jpg") diff --git a/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image.py b/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..5a317eae6268e71004a53c9dfb2c8381fa347568 --- /dev/null +++ b/examples/hidream_o1_image/model_training/validate_full/HiDream-O1-Image.py @@ -0,0 +1,24 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) +state_dict = torch.load("./models/train/HiDream-O1-Image_full/epoch-1.safetensors", map_location="cuda", weights_only=False) +pipe.dit.load_state_dict(state_dict, strict=False) +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=50, +) +image.save("image.jpg") diff --git a/examples/hidream_o1_image/model_training/validate_full/HidreamO1-i2L-v2.py b/examples/hidream_o1_image/model_training/validate_full/HidreamO1-i2L-v2.py new file mode 100644 index 0000000000000000000000000000000000000000..cc9233c89a3407ba6683de5c056cc2f0228fff9f --- /dev/null +++ b/examples/hidream_o1_image/model_training/validate_full/HidreamO1-i2L-v2.py @@ -0,0 +1,31 @@ +from diffsynth.diffusion.template import TemplatePipeline +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig +from modelscope import snapshot_download +from PIL import Image +import numpy as np +import torch +from diffsynth import load_state_dict + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors")], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) +pipe.enable_lora_hot_loading(pipe.dit) +template = TemplatePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ModelConfig(model_id="DiffSynth-Studio/HidreamO1-i2L-v2")], +) +template.models[0].load_state_dict(load_state_dict("models/train/HidreamO1-i2L-v2_full/epoch-1.safetensors")) +snapshot_download("DiffSynth-Studio/HidreamO1-i2L-v2", allow_file_pattern="assets/*", local_dir="data") +images = [Image.open(f"data/assets/multi_input_{i}.jpg") for i in range(4)] +image = template( + pipe, + prompt="A cat is sitting on a stone", + seed=0, cfg_scale=4, num_inference_steps=50, + template_inputs = [{"image": images}], + negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}], +) +image.save("image_output.jpg") \ No newline at end of file diff --git a/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image-Dev.py b/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image-Dev.py new file mode 100644 index 0000000000000000000000000000000000000000..18930ea6ad71b940a6b6f7315e4617c5fb2335b1 --- /dev/null +++ b/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image-Dev.py @@ -0,0 +1,26 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image-Dev", origin_file_pattern="./"), +) +pipe.scheduler = HiDreamO1FlashScheduler(noise_scale_start=7.5, noise_scale_end=7.5, noise_clip_std=2.5) + +pipe.load_lora(pipe.dit, "./models/train/HiDream-O1-Image-Dev_lora/epoch-4.safetensors") +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + cfg_scale=1.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=28, + model_type="dev", + noise_scale=7.5, +) +image.save("image.jpg") diff --git a/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image.py b/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image.py new file mode 100644 index 0000000000000000000000000000000000000000..819a70ab658ea77d4f2c801df3855029bcb5051c --- /dev/null +++ b/examples/hidream_o1_image/model_training/validate_lora/HiDream-O1-Image.py @@ -0,0 +1,23 @@ +import torch +from diffsynth.pipelines.hidream_o1_image import HiDreamO1ImagePipeline, ModelConfig + + +pipe = HiDreamO1ImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="model-*.safetensors"), + ], + processor_config=ModelConfig(model_id="HiDream-ai/HiDream-O1-Image", origin_file_pattern="./"), +) +pipe.load_lora(pipe.dit, "./models/train/HiDream-O1-Image_lora/epoch-4.safetensors") +image = pipe( + prompt="dog,white and brown dog, sitting on wall, under pink flowers", + negative_prompt=" ", + cfg_scale=4.0, + height=2048, + width=2048, + seed=42, + num_inference_steps=50, +) +image.save("image.jpg") diff --git a/examples/ideogram4/model_inference/ideogram-4-bf16-repackage.py b/examples/ideogram4/model_inference/ideogram-4-bf16-repackage.py new file mode 100644 index 0000000000000000000000000000000000000000..0329459e11e433feadf6c734ba6a74ea2097174b --- /dev/null +++ b/examples/ideogram4/model_inference/ideogram-4-bf16-repackage.py @@ -0,0 +1,52 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + + +pipe = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors"), + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="text_encoder/model.safetensors"), + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"), +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [78, 447, 117, 510], "text": "ORACLE\nRed Bull\nRacing", "desc": "Very small white and orange logo on the front of the navy blue cap."}, + {"type": "text", "bbox": [78, 417, 120, 440], "text": "1", "desc": "Bold red numeral '1' on the front left side of the navy blue cap."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [373, 490, 423, 532], "text": "RAUCH", "desc": "Small yellow and blue logo on the right chest of the uniform."}, + {"type": "text", "bbox": [422, 473, 500, 532], "text": "BYBIT\nHONDA", "desc": "Medium-sized white sans-serif text on the right chest of the uniform."}, + {"type": "text", "bbox": [410, 203, 442, 257], "text": "RAUCH", "desc": "Small yellow logo on the left upper arm of the uniform."}, + {"type": "text", "bbox": [530, 448, 627, 510], "text": "Red Bull", "desc": "Medium red text logo on the right side of the torso, part of the Red Bull graphic."}, + {"type": "text", "bbox": [680, 417, 768, 523], "text": "Red Bull", "desc": "Large red text logo across the lower torso of the uniform."}, + {"type": "text", "bbox": [797, 475, 815, 518], "text": "MAX", "desc": "Small white text next to a Dutch flag on the belt area of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."}, + {"type": "text", "bbox": [560, 800, 582, 835], "text": "IA.COM", "desc": "Small blue sans-serif text on the right sleeve of the white shirt."}, + {"type": "text", "bbox": [968, 8, 997, 332], "text": "© Anadolu Agency via Getty Images", "desc": "Small white watermark text in the bottom left corner."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-bf16-repackage.jpg") diff --git a/examples/ideogram4/model_inference/ideogram-4-fp8.py b/examples/ideogram4/model_inference/ideogram-4-fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..3b83e8be03b0ea6c9b8335bfcc2c56c80f19cd61 --- /dev/null +++ b/examples/ideogram4/model_inference/ideogram-4-fp8.py @@ -0,0 +1,51 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + +pipe = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors"), + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="text_encoder/model.safetensors"), + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"), +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [78, 447, 117, 510], "text": "ORACLE\nRed Bull\nRacing", "desc": "Very small white and orange logo on the front of the navy blue cap."}, + {"type": "text", "bbox": [78, 417, 120, 440], "text": "1", "desc": "Bold red numeral '1' on the front left side of the navy blue cap."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [373, 490, 423, 532], "text": "RAUCH", "desc": "Small yellow and blue logo on the right chest of the uniform."}, + {"type": "text", "bbox": [422, 473, 500, 532], "text": "BYBIT\nHONDA", "desc": "Medium-sized white sans-serif text on the right chest of the uniform."}, + {"type": "text", "bbox": [410, 203, 442, 257], "text": "RAUCH", "desc": "Small yellow logo on the left upper arm of the uniform."}, + {"type": "text", "bbox": [530, 448, 627, 510], "text": "Red Bull", "desc": "Medium red text logo on the right side of the torso, part of the Red Bull graphic."}, + {"type": "text", "bbox": [680, 417, 768, 523], "text": "Red Bull", "desc": "Large red text logo across the lower torso of the uniform."}, + {"type": "text", "bbox": [797, 475, 815, 518], "text": "MAX", "desc": "Small white text next to a Dutch flag on the belt area of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."}, + {"type": "text", "bbox": [560, 800, 582, 835], "text": "IA.COM", "desc": "Small blue sans-serif text on the right sleeve of the white shirt."}, + {"type": "text", "bbox": [968, 8, 997, 332], "text": "© Anadolu Agency via Getty Images", "desc": "Small white watermark text in the bottom left corner."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-fp8.jpg") diff --git a/examples/ideogram4/model_inference/ideogram-4-nf4.py b/examples/ideogram4/model_inference/ideogram-4-nf4.py new file mode 100644 index 0000000000000000000000000000000000000000..5986bbd6363c9cb6d1f0221b9a7da3f4fbc3682c --- /dev/null +++ b/examples/ideogram4/model_inference/ideogram-4-nf4.py @@ -0,0 +1,51 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + +pipe = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors"), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors"), + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="text_encoder/model.safetensors"), + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="tokenizer/"), +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [78, 447, 117, 510], "text": "ORACLE\nRed Bull\nRacing", "desc": "Very small white and orange logo on the front of the navy blue cap."}, + {"type": "text", "bbox": [78, 417, 120, 440], "text": "1", "desc": "Bold red numeral '1' on the front left side of the navy blue cap."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [373, 490, 423, 532], "text": "RAUCH", "desc": "Small yellow and blue logo on the right chest of the uniform."}, + {"type": "text", "bbox": [422, 473, 500, 532], "text": "BYBIT\nHONDA", "desc": "Medium-sized white sans-serif text on the right chest of the uniform."}, + {"type": "text", "bbox": [410, 203, 442, 257], "text": "RAUCH", "desc": "Small yellow logo on the left upper arm of the uniform."}, + {"type": "text", "bbox": [530, 448, 627, 510], "text": "Red Bull", "desc": "Medium red text logo on the right side of the torso, part of the Red Bull graphic."}, + {"type": "text", "bbox": [680, 417, 768, 523], "text": "Red Bull", "desc": "Large red text logo across the lower torso of the uniform."}, + {"type": "text", "bbox": [797, 475, 815, 518], "text": "MAX", "desc": "Small white text next to a Dutch flag on the belt area of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."}, + {"type": "text", "bbox": [560, 800, 582, 835], "text": "IA.COM", "desc": "Small blue sans-serif text on the right sleeve of the white shirt."}, + {"type": "text", "bbox": [968, 8, 997, 332], "text": "© Anadolu Agency via Getty Images", "desc": "Small white watermark text in the bottom left corner."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-nf4.jpg") diff --git a/examples/ideogram4/model_inference_low_vram/ideogram-4-bf16-repackage.py b/examples/ideogram4/model_inference_low_vram/ideogram-4-bf16-repackage.py new file mode 100644 index 0000000000000000000000000000000000000000..e7fde1c001f67aedfcf89d9d79a5adc4150603a4 --- /dev/null +++ b/examples/ideogram4/model_inference_low_vram/ideogram-4-bf16-repackage.py @@ -0,0 +1,62 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + +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 = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors", **vram_config), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + # ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="text_encoder/model.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/ideogram-4-bf16-repackage", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [78, 447, 117, 510], "text": "ORACLE\nRed Bull\nRacing", "desc": "Very small white and orange logo on the front of the navy blue cap."}, + {"type": "text", "bbox": [78, 417, 120, 440], "text": "1", "desc": "Bold red numeral '1' on the front left side of the navy blue cap."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [373, 490, 423, 532], "text": "RAUCH", "desc": "Small yellow and blue logo on the right chest of the uniform."}, + {"type": "text", "bbox": [422, 473, 500, 532], "text": "BYBIT\nHONDA", "desc": "Medium-sized white sans-serif text on the right chest of the uniform."}, + {"type": "text", "bbox": [410, 203, 442, 257], "text": "RAUCH", "desc": "Small yellow logo on the left upper arm of the uniform."}, + {"type": "text", "bbox": [530, 448, 627, 510], "text": "Red Bull", "desc": "Medium red text logo on the right side of the torso, part of the Red Bull graphic."}, + {"type": "text", "bbox": [680, 417, 768, 523], "text": "Red Bull", "desc": "Large red text logo across the lower torso of the uniform."}, + {"type": "text", "bbox": [797, 475, 815, 518], "text": "MAX", "desc": "Small white text next to a Dutch flag on the belt area of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."}, + {"type": "text", "bbox": [560, 800, 582, 835], "text": "IA.COM", "desc": "Small blue sans-serif text on the right sleeve of the white shirt."}, + {"type": "text", "bbox": [968, 8, 997, 332], "text": "© Anadolu Agency via Getty Images", "desc": "Small white watermark text in the bottom left corner."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-bf16-repackage.jpg") diff --git a/examples/ideogram4/model_inference_low_vram/ideogram-4-fp8.py b/examples/ideogram4/model_inference_low_vram/ideogram-4-fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..59cfa2dc100c72fec60c70255632b90471c2fd56 --- /dev/null +++ b/examples/ideogram4/model_inference_low_vram/ideogram-4-fp8.py @@ -0,0 +1,52 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.bfloat16, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors", **vram_config), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors", **vram_config), + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="text_encoder/model.safetensors", **vram_config), + ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-fp8", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-fp8_low_vram.jpg") diff --git a/examples/ideogram4/model_inference_low_vram/ideogram-4-nf4.py b/examples/ideogram4/model_inference_low_vram/ideogram-4-nf4.py new file mode 100644 index 0000000000000000000000000000000000000000..2634a0716f8a7b207ec21f8422f21e126e7e7583 --- /dev/null +++ b/examples/ideogram4/model_inference_low_vram/ideogram-4-nf4.py @@ -0,0 +1,52 @@ +from diffsynth.pipelines.ideogram4 import Ideogram4Pipeline +from diffsynth.core import ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.bfloat16, + "onload_device": "cpu", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = Ideogram4Pipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="transformer/diffusion_pytorch_model.safetensors", **vram_config), + # unconditional_transformer is optional. You can delete this line to reduce VRAM required. + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="unconditional_transformer/diffusion_pytorch_model.safetensors", **vram_config), + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="text_encoder/model.safetensors", **vram_config), + ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="ideogram-ai/ideogram-4-nf4", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = r""" +{ + "high_level_description": "A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull Racing racing suit and cap, smiling as he holds his racing helmet and talks to a man in a white shirt and black vest at a race track.", + "style_description": { + "aesthetics": "saturated primary colors, rule of thirds, joyful and triumphant", + "lighting": "overcast daylight, diffused, soft subtle shadows", + "photo": "shallow depth of field, sharp focus, eye-level, telephoto", + "medium": "photograph" + }, + "compositional_deconstruction": { + "background": "The background is an out-of-focus racing paddock or track environment. Several blurred figures are visible, including one in an orange shirt. A purple and white structure with a red 'F1' logo stands on the left. The scene is outdoors with daylight, though the sky is not visible.", + "elements": [ + {"type": "obj", "bbox": [55, 642, 1000, 937], "desc": "An older man standing in profile, facing left toward Max Verstappen. He has grey hair and fair skin. He is wearing a white long-sleeved button-down shirt with a navy blue quilted vest over it. He has a slight smile."}, + {"type": "obj", "bbox": [34, 137, 1000, 617], "desc": "Max Verstappen, a fair-skinned male Formula 1 driver, positioned in the center. He is facing forward with a joyful expression and a slight smile. He wears a navy blue Red Bull Racing team uniform with numerous sponsor logos and a matching baseball cap with the number '1'. He is holding a white and red racing helmet in his hands. He has a silver watch on his left wrist."}, + {"type": "obj", "bbox": [422, 212, 792, 452], "desc": "Max Verstappen's racing helmet, held in front of his chest. It features a white, red, and yellow design with the Red Bull logo and the 'Player 0.0' branding. The visor is clear and open."}, + {"type": "text", "bbox": [657, 0, 755, 142], "text": "F1", "desc": "Large, stylized red logo on a black and purple background in the lower left."}, + {"type": "text", "bbox": [768, 0, 818, 147], "text": "Formula 1\nWorld Championship™", "desc": "Small white sans-serif text below the F1 logo on the left side."}, + {"type": "text", "bbox": [332, 442, 363, 483], "text": "Red Bull", "desc": "Small yellow and red text logo on the collar of the uniform."}, + {"type": "text", "bbox": [558, 317, 715, 355], "text": "Player 0.0", "desc": "Black sans-serif text on a white band on the racing helmet."} + ] + } +} +""" +image = pipe(prompt=prompt, height=1024, width=1024, num_inference_steps=48, cfg_scale=7.0, seed=42) +image.save("image_ideogram-4-nf4_low_vram.jpg") diff --git a/examples/ideogram4/model_training/full/Ideogram-4-bf16-repackage.sh b/examples/ideogram4/model_training/full/Ideogram-4-bf16-repackage.sh new file mode 100644 index 0000000000000000000000000000000000000000..a27c4b7c816addac62639497a11a9fa418474b3c --- /dev/null +++ b/examples/ideogram4/model_training/full/Ideogram-4-bf16-repackage.sh @@ -0,0 +1,17 @@ +# Dataset: data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage/ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ideogram4/Ideogram-4-bf16-repackage/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch --config_file examples/ideogram4/model_training/full/accelerate_config.yaml examples/ideogram4/model_training/train.py \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --dataset_base_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage" \ + --dataset_metadata_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage/metadata.json" \ + --model_id_with_origin_paths "DiffSynth-Studio/ideogram-4-bf16-repackage:transformer/diffusion_pytorch_model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:text_encoder/model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:vae/diffusion_pytorch_model.safetensors" \ + --remove_prefix_in_ckpt "pipe.dit." \ + --trainable_models "dit" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --output_path "./models/train/Ideogram-4-bf16-repackage_full" \ + --data_file_keys "image" diff --git a/examples/ideogram4/model_training/full/accelerate_config.yaml b/examples/ideogram4/model_training/full/accelerate_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..83280f73f315a32eccb065f351d66b4b2678759d --- /dev/null +++ b/examples/ideogram4/model_training/full/accelerate_config.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: false + zero_stage: 2 +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 diff --git a/examples/ideogram4/model_training/lora/Ideogram-4-bf16-repackage.sh b/examples/ideogram4/model_training/lora/Ideogram-4-bf16-repackage.sh new file mode 100644 index 0000000000000000000000000000000000000000..ac4153b3621b9e1b8e60dcba1ce7890fad0859b1 --- /dev/null +++ b/examples/ideogram4/model_training/lora/Ideogram-4-bf16-repackage.sh @@ -0,0 +1,19 @@ +# Dataset: data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage/ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ideogram4/Ideogram-4-bf16-repackage/*" --local_dir ./data/diffsynth_example_dataset + +accelerate launch examples/ideogram4/model_training/train.py \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --dataset_base_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage" \ + --dataset_metadata_path "./data/diffsynth_example_dataset/ideogram4/Ideogram-4-bf16-repackage/metadata.json" \ + --model_id_with_origin_paths "DiffSynth-Studio/ideogram-4-bf16-repackage:transformer/diffusion_pytorch_model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:text_encoder/model.safetensors,DiffSynth-Studio/ideogram-4-bf16-repackage:vae/diffusion_pytorch_model.safetensors" \ + --lora_base_model "dit" \ + --remove_prefix_in_ckpt "pipe.dit." \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --output_path "./models/train/Ideogram-4-bf16-repackage_lora" \ + --lora_target_modules "attention.qkv,attention.o,feed_forward.w1,feed_forward.w2,feed_forward.w3,adaln_modulation" \ + --data_file_keys "image"