DiffSynth-Studio Documentation

Technical Documentation: modelscope/DiffSynth-Studio

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (`README.md`) Copy Section

DiffSynth-Studio

切换到中文版

Introduction Copy Chapter

Welcome to the magical world of Diffusion models! DiffSynth-Studio is an open-source Diffusion model engine developed and maintained by the ModelScope Community team. We hope to foster technological innovation through framework construction, aggregate the power of the open-source community, and explore the rich capabilities of generative model technology!

Framework features:

• Model Support: The framework integrates mainstream open-source Diffusion models, covering image generation, video generation, audio generation, and image quality metrics models.
• VRAM Management: Dynamically schedules model parameters across disk, memory, and VRAM, allowing consumer-grade GPUs with low VRAM to run inference with large models.
• Parameter Quantization: Converts model parameters to quantized precisions such as NF4 and INT8, significantly reducing the VRAM requirements of model inference and LoRA training.
• Arbitrary Training: Almost every model that supports inference also supports training, whether base models, LoRAs, or any Adapter models with additional inputs.
• Split Training: Uses a computational graph inference engine to track every variable in the Pipeline, splitting the training process into two stages for efficient training.

References:

• Developer documentation (for humans): 中文版、English version
• Agent Skills (for AI): DiffSynth-Studio Model Integration Skills

See also:

• DiffSynth-ComfyUI: A ComfyUI node package built on DiffSynth-Studio, enabling you to build workflows in ComfyUI to run model inference.
• DiffSynth-WebUI: A lightweight LoRA training tool built on DiffSynth-Studio, enabling LoRA training of models on consumer-grade GPUs.
• ModelScope AIGC Zone (for Chinese users): Productized features powered by DiffSynth-Studio as the core inference and training engine; experience the rich potential of the open-source model ecosystem.
• ModelScope Civision (for global users): Unlock the vast potential of the open-source model ecosystem through productized capabilities powered by DiffSynth-Studio.

Update History Copy Chapter

DiffSynth-Studio has undergone major version updates, and some old features are no longer maintained. If you need to use old features, please switch to the last historical version before the major version update.

Currently, the development personnel of this project are limited, with most of the work handled by Artiprocher and mi804. Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand.

• October 21, 2026 We now support Fibo 2.0, including unified text-to-image and image editing inference, masked and multi-reference editing, text rendering, low VRAM inference, and full/LoRA training for both text-to-image and image editing on Fibo-2.0-dev. See the documentation and example code.

• October 20, 2026 We now support MOSS-Image, including inference and training for MOSS-Image-Base and MOSS-Image-SFT. In addition, we provide an EntroPack low-bit quantized model, which can run the 3.5-bit quantized model with as little as 4 GB of VRAM; and an accelerated model distillation-accelerated via OPD training, supporting 8-step ultra-fast inference.

• October 16, 2026 flux-3-image-klein open-sourced, welcome a new member to the image model family! Support includes text-to-image generation, multi-reference editing, layout control, low VRAM inference, and Base training capabilities. For details, please refer to the documentation and example code. We quantized the model with EntroPack, providing pre-quantized models in 3-bit, 4-bit, and 6-bit versions at DiffSynth-Studio/flux-3-image-klein-entropack.

• October 12, 2026 We now support distillation accelerated training based on OPD (On-Policy Distillation), which enables distillation acceleration of models at an extremely low cost. For details, please refer to Distillation Accelerated Training.

• September 29, 2026 We are excited to open-source EntroPack, a general-purpose tensor compression library for PyTorch that provides lossless compression with exact recovery and rate-controlled lossy compression, with GPU encoding and decoding. It is integrated into DiffSynth-Studio as a quantization backend, compressing model weights at arbitrary target bitrates in both inference and LoRA training. For more information, please refer to:

  • Project Page: https://github.com/modelscope/entropack
  • Technical Report: https://arxiv.org/abs/2609.34185
  • Pre-quantized Models: https://www.modelscope.cn/collections/DiffSynth-Studio/EntroPack-Quantized-Models
  • Example Code: Z-Image-Turbo, Qwen-Image-2.1, MiniMax-H3, LoRA Training Example

• September 20, 2026 Qwen-Image-2.1 was open-sourced. The model supports text-to-image generation, image editing, and RGBA output with a transparency channel. DiffSynth-Studio now provides full support for it. For details, please refer to the documentation and example code.

• September 14, 2026 We trained and open-sourced DiffSynth-Music, a controllable music generation model based on ACE-Step, which supports five control modes: Beats, Vocals, Accompany, Prosody, and Reference. For more information, please refer to:

  • Model: https://modelscope.cn/models/DiffSynth-Studio/DiffSynth-Music
  • Technical Report: https://arxiv.org/abs/2609.12774
  • Code: ./examples/diffsynth_music/

• September 11, 2026 We have added support for YuE2, a new music generation model. For details, please refer to the documentation and example code.

• September 10, 2026 We have integrated LTX-2.5, the latest audio-video generation model from Lightricks. The features include text-to-audio/video with automatic duration prediction, image-to-audio/video with keyframe interpolation, audio-to-video, audio-video retake, IC-LoRA pixel spatial upscaling, text-to-audio, INT8 quantized inference, low VRAM inference, and training. For details, please refer to the documentation and code.

• September 10, 2026 We have open-sourced DiffSynth-ComfyUI. Now you can run model inference with DiffSynth-Studio in ComfyUI.

• September 8, 2026 We trained a MiniMax-H3 Training Adapter for the CFG-distilled MiniMax-H3 base, a DeCFG LoRA training adapter offered in FL2VA and Ref2VA versions. The training data comes from the self-generated MiniMax-H3-Self-Generated-Dataset. Based on this adapter, we trained two toy models, MiniMax-H3-Songyu-LoRA (character identity, FL2VA) and MiniMax-H3-Ref2VA-FirstFrame-Lineart (lineart first-frame control, Ref2VA), as references for inference and fine-tuning.

• September 1, 2026 We have integrated SenseNova-U1.5, SenseTime's unified multimodal model, for which we provide text-to-image generation, image editing, low VRAM inference, and training support. For details, please refer to the documentation and example code.

• August 31, 2026 We have integrated Qwen-Video-Edit, a video editing model developed by open-source community contributor yunpeng1998 based on the image editing model Qwen-Image-Edit. This serves as an excellent example of exploring and expanding model capabilities.

• August 25, 2026 We have open-sourced DiffSynth-WebUI, enabling one-click private deployment of LoRA training services. Combined with the model quantization feature, you can train large models even with consumer-grade GPUs.

• August 19, 2026 We have released the model quantization feature. It provides a unified QuantizeConfig entry point supporting multiple quantization backends including bitsandbytes, torchao, and comfy-kitchen, with capabilities such as online quantization, loading pre-quantized weights, mixed quantization, saving quantized models, and quantization + LoRA training. For details, please refer to the documentation.

More

• August 17, 2026 MiniMax-Music3 open-sourced, welcome a new member to the audio model family! Support includes text-to-music generation and low VRAM inference. For details, please refer to the documentation and example code.

• August 7, 2026 We add support for Wan-Animate-2 in the Wan series. Given a reference image and a driving video, it makes the reference character perform the motions in the driving video, generating high-quality character animation, with both standard and distilled variants. For details, please refer to the documentation and example code.

• August 3, 2026 MiniMax-H3 open-sourced, welcome a new member to the video model family! Support includes text-to-video-audio generation, keyframe-guided generation, reference-driven generation, low VRAM inference, and NF4-quantized inference. For details, please refer to the documentation and example code.

• July 28, 2026 LingBot-Video open-sourced, welcome a new member to the video model family! This release includes two variants, Dense-1.3B and MoE-30B-A3B (30B total parameters, ~3B active per token), both supporting text-to-video, image-to-video and text-to-image generation, low VRAM inference, and LoRA / full training capabilities. For details, please refer to the documentation and example code. Huge thanks to NancyFyong for contributing the integration of this model!

• July 21, 2026 We have open-sourced DiffSynth-Studio Model Integration Skills. This is a composable collection of Agent Skills that automates the entire workflow of integrating external diffusion models into DiffSynth-Studio, significantly improving the standardization and efficiency of model integration. Get started with the example!

• June 29, 2026 Boogu-Image open-sourced. Support includes text-to-image generation, image editing, low VRAM inference, and training capabilities. For details, please refer to the documentation and example code.

• June 24, 2026 Krea-2 is now open-source, and we have provided full support. For more details, please refer to the documentation and example code.

• June 16, 2026: We have added a new Template model for ACE-Step: vocals2music. For more details, please refer to the documentation and example code.

• June 15, 2026 We have open-sourced Image-to-LoRA V2, compressing the hours-long training process for image style LoRAs into a single model inference step, thereby exploring a new paradigm for LoRA model training. The technical report has been released. This release includes three models:

  • DiffSynth-Studio/ZImage-i2L-v2: Adapted for the Z-Image model
  • DiffSynth-Studio/KleinBase4B-i2L-v2: Adapted for the FLUX.2-klein-base-4B model
  • DiffSynth-Studio/HidreamO1-i2L-v2: Adapted for the Hidream-O1-Image model

• June 5, 2026 Ideogram 4 open-sourced. Support includes text-to-image inference. For details, please refer to the documentation and example code.

• May 21, 2026: Added support for image quality metrics models, including FID, CLIP, Aesthetic, PickScore, ImageReward, HPSv2, and HPSv3. For details, refer to the documentation and example code.

• May 18, 2026 Added CPU Offload Training support. By moving model weights layer-by-layer between CPU and GPU, it significantly reduces GPU VRAM usage during training, enabling LoRA training of large models even on consumer-grade GPUs, compatible with all models. Simply add --enable_model_cpu_offload to your training command to enable (currently supports single-GPU training only). For details, see the documentation.

• May 14, 2026 HiDream-O1-Image open-sourced, welcome a new member to the image model family! Support includes text-to-image generation, image editing, low VRAM inference, and training capabilities. For details, please refer to the documentation and example code.

• April 28, 2026 We released Diffusion Templates, a plugin framework designed for Diffusion models that significantly lowers the barrier to training controllable generative models. Let's explore this cutting-edge technology together!

  • Open-source code: DiffSynth-Studio
  • Technical report: arXiv
  • Project homepage: GitHub
  • Documentation: English Version | Chinese Version
  • Online demo: ModelScope
  • Model collections: ModelScope | ModelScope International | HuggingFace
  • Datasets: ModelScope | ModelScope International | HuggingFace

• April 27, 2026 We support ACE-Step-1.5! Support includes text-to-music generation, low VRAM inference, and LoRA training capabilities. For details, please refer to the documentation and example code.

• April 27, 2026: We have reinstated support for the Stable Diffusion v1.5 and SDXL models, providing academic research support exclusively for these two model types.

• April 14, 2026 JoyAI-Image open-sourced, welcome a new member to the image editing model family! Support includes instruction-guided image editing, low VRAM inference, and training capabilities. For details, please refer to the documentation and example code.

• March 19, 2026: Added support for openmoss/MOVA-720p and openmoss/MOVA-360p models, including training and inference capabilities. Documentation and example code are now available.

• March 12, 2026: We have added support for the LTX-2.3 audio-video generation model. The features includes text-to-audio/video, image-to-audio/video, IC-LoRA control, audio-to-video, and audio-video inpainting. We have supported the complete inference and training functionalities. For details, please refer to the documentation and code.

• March 3, 2026: We released the DiffSynth-Studio/Qwen-Image-Layered-Control-V2 model, which is an updated version of Qwen-Image-Layered-Control. In addition to the originally supported text-guided functionality, it adds brush-controlled layer separation capabilities.

• March 2, 2026 Added support for Anima. For details, please refer to the documentation. This is an interesting anime-style image generation model. We look forward to its future updates.

• February 26, 2026 Added full and lora training support for the LTX-2 audio-video generation model. See the documentation for details.

• February 10, 2026 Added inference support for the LTX-2 audio-video generation model. See the documentation for details. Support for model training will be implemented in the future.

• February 2, 2026 The first document of the Research Tutorial series is now available, guiding you through training a small 0.1B text-to-image model from scratch. For details, see the documentation and model. We hope DiffSynth-Studio can evolve into a more powerful training framework for Diffusion models.

• January 27, 2026: Z-Image is released, and our Z-Image-i2L model is released concurrently. You can use it in ModelScope Studios. For details, see the documentation.

• January 19, 2026: Added support for FLUX.2-klein-4B and FLUX.2-klein-9B models, including training and inference capabilities. Documentation and example code are now available.

• January 12, 2026: We trained and open-sourced a text-guided image layer separation model (Model Link). Given an input image and a textual description, the model isolates the image layer corresponding to the described content. For more details, please refer to our blog post (Chinese version, English version).

• December 24, 2025: Based on Qwen-Image-Edit-2511, we trained an In-Context Editing LoRA model (Model Link). This model takes three images as input (Image A, Image B, and Image C), and automatically analyzes the transformation from Image A to Image B, then applies the same transformation to Image C to generate Image D. For more details, please refer to our blog post (Chinese version, English version).

• December 9, 2025 We release a wild model based on DiffSynth-Studio 2.0: Qwen-Image-i2L (Image-to-LoRA). This model takes an image as input and outputs a LoRA. Although this version still has significant room for improvement in terms of generalization, detail preservation, and other aspects, we are open-sourcing these models to inspire more innovative research. For more details, please refer to our blog.

• December 4, 2025 DiffSynth-Studio 2.0 released! Many new features online

  • Documentation online: Our documentation is still continuously being optimized and updated
  • VRAM Management module upgraded, supporting layer-level disk offload, releasing both memory and VRAM simultaneously
  • New model support
    • Z-Image Turbo: Model, Documentation, Code
    • FLUX.2-dev: Model, Documentation, Code
  • Training framework upgrade
    • Split Training: Supports automatically splitting the training process into two stages: data processing and training (even for training ControlNet or any other model). Computations that do not require gradient backpropagation, such as text encoding and VAE encoding, are performed during the data processing stage, while other computations are handled during the training stage. Faster speed, less VRAM requirement.
    • Differential LoRA Training: This is a training technique we used in ArtAug, now available for LoRA training of any model.
    • FP8 Training: FP8 can be applied to any non-training model during training, i.e., models with gradients turned off or gradients that only affect LoRA weights.

• November 4, 2025 Supported the ByteDance/Video-As-Prompt-Wan2.1-14B model, which is trained based on Wan 2.1 and supports generating corresponding actions based on reference videos.

• October 30, 2025 Supported the meituan-longcat/LongCat-Video model, which supports text-to-video, image-to-video, and video continuation. This model uses the Wan framework for inference and training in this project.

• October 27, 2025 Supported the krea/krea-realtime-video model, adding another member to the Wan model ecosystem.

• September 23, 2025 DiffSynth-Studio/Qwen-Image-EliGen-Poster released! This model was jointly developed and open-sourced by us and Taobao Experience Design Team. Built upon Qwen-Image, the model is specifically designed for e-commerce poster scenarios, supporting precise partition layout control. Please refer to our sample code.

• September 9, 2025 Our training framework supports various training modes. Currently adapted for Qwen-Image, in addition to the standard SFT training mode, Direct Distill is now supported. Please refer to our sample code. This feature is experimental, and we will continue to improve it to support more comprehensive model training functions.

• August 28, 2025 We support Wan2.2-S2V, an audio-driven cinematic video generation model. See ./examples/wanvideo/.

• August 21, 2025 DiffSynth-Studio/Qwen-Image-EliGen-V2 released! Compared to the V1 version, the training dataset has been changed to Qwen-Image-Self-Generated-Dataset, so the generated images better conform to Qwen-Image's own image distribution and style. Please refer to our sample code.

• August 21, 2025 We open-sourced the DiffSynth-Studio/Qwen-Image-In-Context-Control-Union structural control LoRA model, adopting the In Context technical route, supporting multiple categories of structural control conditions, including canny, depth, lineart, softedge, normal, and openpose. Please refer to our sample code.

• August 20, 2025 We open-sourced the DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix model, improving the editing effect of Qwen-Image-Edit on low-resolution image inputs. Please refer to our sample code

• August 19, 2025 Qwen-Image-Edit open-sourced, welcome a new member to the image editing model family!

• August 18, 2025 We trained and open-sourced the Qwen-Image inpainting ControlNet model DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint. The model structure adopts a lightweight design. Please refer to our sample code.

• August 15, 2025 We open-sourced the Qwen-Image-Self-Generated-Dataset dataset. This is an image dataset generated using the Qwen-Image model, containing 160,000 1024 x 1024 images. It includes general, English text rendering, and Chinese text rendering subsets. We provide annotations for image descriptions, entities, and structural control images for each image. Developers can use this dataset to train Qwen-Image models' ControlNet and EliGen models. We aim to promote technological development through open-sourcing!

• August 13, 2025 We trained and open-sourced the Qwen-Image ControlNet model DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Depth. The model structure adopts a lightweight design. Please refer to our sample code.

• August 12, 2025 We trained and open-sourced the Qwen-Image ControlNet model DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny. The model structure adopts a lightweight design. Please refer to our sample code.

• August 11, 2025 We open-sourced the distilled acceleration model DiffSynth-Studio/Qwen-Image-Distill-LoRA for Qwen-Image, following the same training process as DiffSynth-Studio/Qwen-Image-Distill-Full, but the model structure has been modified to LoRA, thus being better compatible with other open-source ecosystem models.

• August 7, 2025 We open-sourced the entity control LoRA model DiffSynth-Studio/Qwen-Image-EliGen for Qwen-Image. Qwen-Image-EliGen can achieve entity-level controlled text-to-image generation. Technical details can be found in the paper. Training dataset: EliGenTrainSet.

• August 5, 2025 We open-sourced the distilled acceleration model DiffSynth-Studio/Qwen-Image-Distill-Full for Qwen-Image, achieving approximately 5x acceleration.

• August 4, 2025 Qwen-Image open-sourced, welcome a new member to the image generation model family!

• August 1, 2025 FLUX.1-Krea-dev open-sourced, a text-to-image model focused on aesthetic photography. We provided comprehensive support in a timely manner, including low VRAM layer-by-layer offload, LoRA training, and full training. For more details, please refer to ./examples/flux/.

• July 28, 2025 Wan 2.2 open-sourced. We provided comprehensive support in a timely manner, including low VRAM layer-by-layer offload, FP8 quantization, sequence parallelism, LoRA training, and full training. For more details, please refer to ./examples/wanvideo/.

• July 11, 2025 We propose Nexus-Gen, a unified framework that combines the language reasoning capabilities of Large Language Models (LLMs) with the image generation capabilities of diffusion models. This framework supports seamless image understanding, generation, and editing tasks.

  • Paper: Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space
  • GitHub Repository: https://github.com/modelscope/Nexus-Gen
  • Model: ModelScope, HuggingFace
  • Training Dataset: ModelScope Dataset
  • Online Experience: ModelScope Nexus-Gen Studio

• June 15, 2025 ModelScope's official evaluation framework EvalScope now supports text-to-image generation evaluation. Please refer to the best practices guide to try it out.

• March 25, 2025 Our new open-source project DiffSynth-Engine is now open-sourced! Focused on stable model deployment, targeting industry, providing better engineering support, higher computational performance, and more stable features.

• March 31, 2025 We support InfiniteYou, a face feature preservation method for FLUX. More details can be found in ./examples/InfiniteYou/.

• March 13, 2025 We support HunyuanVideo-I2V, the image-to-video generation version of Tencent's open-source HunyuanVideo. More details can be found in ./examples/HunyuanVideo/.

• February 25, 2025 We support Wan-Video, a series of state-of-the-art video synthesis models open-sourced by Alibaba. See ./examples/wanvideo/.

• February 17, 2025 We support StepVideo! Advanced video synthesis model! See ./examples/stepvideo.

• December 31, 2024 We propose EliGen, a new framework for entity-level controlled text-to-image generation, supplemented with an inpainting fusion pipeline, extending its capabilities to image inpainting tasks. EliGen can seamlessly integrate existing community models such as IP-Adapter and In-Context LoRA, enhancing their versatility. For more details, see ./examples/EntityControl.

  • Paper: EliGen: Entity-Level Controlled Image Generation with Regional Attention
  • Model: ModelScope, HuggingFace
  • Online Experience: ModelScope EliGen Studio
  • Training Dataset: EliGen Train Set

• December 19, 2024 We implemented advanced VRAM management for HunyuanVideo, enabling video generation with resolutions of 129x720x1280 on 24GB VRAM or 129x512x384 on just 6GB VRAM. More details can be found in ./examples/HunyuanVideo/.

• December 18, 2024 We propose ArtAug, a method to improve text-to-image models through synthesis-understanding interaction. We trained an ArtAug enhancement module for FLUX.1-dev in LoRA format. This model incorporates the aesthetic understanding of Qwen2-VL-72B into FLUX.1-dev, thereby improving the quality of generated images.

  • Paper: https://arxiv.org/abs/2412.12888
  • Example: https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/ArtAug
  • Model: ModelScope, HuggingFace
  • Demo: ModelScope, HuggingFace (coming soon)

• October 25, 2024 We provide extensive FLUX ControlNet support. This project supports many different ControlNet models and can be freely combined, even if their structures are different. Additionally, ControlNet models are compatible with high-resolution optimization and partition control technologies, enabling very powerful controllable image generation. See ./examples/ControlNet/.

• October 8, 2024 We released extended LoRAs based on CogVideoX-5B and ExVideo. You can download this model from ModelScope or HuggingFace.

• August 22, 2024 This project now supports CogVideoX-5B. See here. We provide several interesting features for this text-to-video model, including:

  • Text-to-video
  • Video editing
  • Self super-resolution
  • Video interpolation

• August 22, 2024 We implemented an interesting brush feature that supports all text-to-image models. Now you can create stunning images with the assistance of AI using the brush!

  • Use it in our WebUI.

• August 21, 2024 DiffSynth-Studio now supports FLUX.

  • Enable CFG and high-resolution inpainting to improve visual quality. See here
  • LoRA, ControlNet, and other addon models will be released soon.

• June 21, 2024 We propose ExVideo, a post-training fine-tuning technique aimed at enhancing the capabilities of video generation models. We extended Stable Video Diffusion to achieve long video generation of up to 128 frames.

  • Project Page
  • Source code has been released in this repository. See examples/ExVideo.
  • Model has been released at HuggingFace and ModelScope.
  • Technical report has been released at arXiv.
  • You can try ExVideo in this demo!

• June 13, 2024 DiffSynth Studio has migrated to ModelScope. The development team has also transitioned from "me" to "us". Of course, I will still participate in subsequent development and maintenance work.

• January 29, 2024 We propose Diffutoon, an excellent cartoon coloring solution.

  • Project Page
  • Source code has been released in this project.
  • Technical report (IJCAI 2024) has been released at arXiv.

• December 8, 2023 We decided to initiate a new project aimed at unleashing the potential of diffusion models, especially in video synthesis. The development work of this project officially began.

• November 15, 2023 We propose FastBlend, a powerful video deflickering algorithm.

  • sd-webui extension has been released at GitHub.
  • Demonstration videos have been showcased on Bilibili, including three tasks:
    • Video Deflickering
    • Video Interpolation
    • Image-Driven Video Rendering
  • Technical report has been released at arXiv.
  • Unofficial ComfyUI extensions developed by other users have been released at GitHub.

• October 1, 2023 We released an early version of the project named FastSDXL. This was an initial attempt to build a diffusion engine.

  • Source code has been released at GitHub.
  • FastSDXL includes a trainable OLSS scheduler to improve efficiency.
    • The original repository of OLSS is located here.
    • Technical report (CIKM 2023) has been released at arXiv.
    • Demonstration video has been released at Bilibili.
    • Since OLSS requires additional training, we did not implement it in this project.

• August 29, 2023 We propose DiffSynth, a video synthesis framework.

  • Project Page.
  • Source code has been released at EasyNLP.
  • Technical report (ECML PKDD 2024) has been released at arXiv.

Installation Copy Chapter

Install from source (recommended):

git clone https://github.com/modelscope/DiffSynth-Studio.git  
cd DiffSynth-Studio
pip install -e .

For more installation methods and instructions for non-NVIDIA GPUs, please refer to the Installation Guide.

Download Source Configuration

Before model inference and training, you can configure the model download source and other options through environment variables.

This project downloads models from ModelScope by default. For users outside China, you can download models from the ModelScope International site with the following configuration:

shell Copy

export MODELSCOPE_ENDPOINT=https://modelscope.ai

To download models from HuggingFace, please modify the environment variables. Note that model IDs may differ across platforms:

shell Copy

export DIFFSYNTH_DOWNLOAD_SOURCE="huggingface"

Basic Framework Copy Chapter

DiffSynth-Studio redesigns the inference and training pipelines for mainstream Diffusion models (including FLUX, Wan, etc.), enabling efficient memory management and flexible model training.

Quick start: experience popular and the latest models:

Model overview:

• Image generation
  • Fibo 2.0: Documentation, Example code
  • MOSS-Image: Documentation, Example code
  • FLUX.3: Documentation, Example code
  • Qwen-Image-2.1: Documentation, Example code
  • SenseNova-U1: Documentation, Example code
  • Boogu-Image: Documentation, Example code
  • Krea-2: Documentation, Example code
  • Ideogram 4: Documentation, Example code
  • HiDream-O1-Image: Documentation, Example code
  • JoyAI-Image: Documentation, Example code
  • ERNIE-Image: Documentation, Example code
  • FLUX.2: Documentation, Example code
  • Z-Image: Documentation, Example code
  • Anima: Documentation, Example code
  • Qwen-Image: Documentation, Example code
  • FLUX.1: Documentation, Example code
  • Stable Diffusion XL: Documentation, Example code
  • Stable Diffusion: Documentation, Example code
• Video generation
  • MiniMax-H3: Documentation, Example code
  • LingBot-Video: Documentation, Example code
  • LTX-2: Documentation, Example code
  • Wan: Documentation, Example code
• Audio generation
  • DiffSynth-Music: Documentation, Example code
  • MiniMax-Music3: Documentation, Example code
  • YuE2: Documentation, Example code
  • ACE-Step: Documentation, Example code
• Evaluation models: Documentation, Example code

View all supported models

Innovative Achievements Copy Chapter

We believe that a well-developed open-source code framework can lower the threshold for technical exploration. We have achieved many interesting technologies based on this codebase. Perhaps you also have many wild ideas, and with DiffSynth-Studio, you can quickly realize these ideas.

EntroPack: A General-Purpose Tensor Compression Library for PyTorch

Compress model weights at arbitrary target bitrates with GPU encoding and decoding, and run the compressed weights directly in inference and LoRA training.

• Project Page: GitHub
• Technical Report: https://arxiv.org/abs/2609.34185
• Pre-quantized Models: ModelScope
• Documentation: English Version | Chinese Version
• Example Code: Z-Image-Turbo, Qwen-Image-2.1, MiniMax-H3, LoRA Training Example

TreeAdapter: A Model System Built from Structured LoRAs

A model system composed of 10,000+ LoRAs challenges the precise generation of tens of thousands of rare species.

• Paper: TreeAdapter: Hierarchical Taxonomy-Guided Adapter Composition for Fine-Grained Species Image Generation
• Model: ModelScope

Image-to-LoRA: Compressing Model Training into Model Inference

True Meta Learning: feed a dataset of images into one end of the model, and the trained LoRA model comes out the other end.

• Paper: Compressing Image Style Training into a Single Model Forward
• Model:
  • DiffSynth-Studio/ZImage-i2L-v2: Adapted for the Z-Image model
  • DiffSynth-Studio/KleinBase4B-i2L-v2: Adapted for the FLUX.2-klein-base-4B model
  • DiffSynth-Studio/HidreamO1-i2L-v2: Adapted for the Hidream-O1-Image model

Diffusion-Templates: A Plugin-Based Controllable Generation Framework

One framework that turns every controllable generation capability into a plugin, allowing multiple models to combine and emerge with rich generation capabilities.

• Paper: Diffusion Templates: A Unified Plugin Framework for Controllable Diffusion
• Project homepage: GitHub
• Documentation: English Version | Chinese Version
• Online demo: ModelScope
• Model collections: ModelScope | ModelScope International | HuggingFace
• Datasets: ModelScope | ModelScope International | HuggingFace

Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation

Trade inference time for higher quality of generated content.

• Paper: Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation
• Sample Code: /docs/en/Research_Tutorial/inference_time_scaling.md

VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers

Transform image 3 into image 4 based on the change from image 1 to image 2 — the emergent capability of image editing models.

• Paper: VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers
• Sample code: /examples/qwen_image/model_inference/Qwen-Image-Edit-2511-ICEdit.py
• Model: ModelScope

AttriCtrl: Attribute Intensity Control for Image Generation Models

Numerical attributes can also precisely control image generation models.

• Paper: AttriCtrl: Fine-Grained Control of Aesthetic Attribute Intensity in Diffusion Models
• Sample Code: /examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py
• Model: ModelScope

AutoLoRA: Automated LoRA Retrieval and Fusion

LoRA is a product that unifies needs and solutions — how can we make better use of these LoRAs?

• Paper: AutoLoRA: Automatic LoRA Retrieval and Fine-Grained Gated Fusion for Text-to-Image Generation
• Sample Code: /examples/flux/model_inference/FLUX.1-dev-LoRA-Fusion.py
• Model: ModelScope

Nexus-Gen: Unified Architecture for Image Understanding, Generation, and Editing

What happens when a single model combines image understanding, generation, and editing capabilities?

• Detailed Page: https://github.com/modelscope/Nexus-Gen
• Paper: Nexus-Gen: Unified Image Understanding, Generation, and Editing via Prefilled Autoregression in Shared Embedding Space
• Model: ModelScope, HuggingFace
• Dataset: ModelScope Dataset
• Online Experience: ModelScope Nexus-Gen Studio

ArtAug: Aesthetic Enhancement for Image Generation Models

A single LoRA that significantly enhances detail and aesthetics.

• Detailed Page: ./examples/ArtAug/
• Paper: ArtAug: Enhancing Text-to-Image Generation through Synthesis-Understanding Interaction
• Model: ModelScope, HuggingFace
• Online Experience: ModelScope AIGC Tab

EliGen: Precise Image Partition Control

How can region-based layers control the position of content in an image?

• Paper: EliGen: Entity-Level Controlled Image Generation with Regional Attention
• Sample Code: /examples/flux/model_inference/FLUX.1-dev-EliGen.py
• Model: ModelScope, HuggingFace
• Online Experience: ModelScope EliGen Studio
• Dataset: EliGen Train Set

ExVideo: Extended Training for Video Generation Models

If a video generation model can only generate 25 frames, how can we make it generate longer videos?

• Project Page: Project Page
• Paper: ExVideo: Extending Video Diffusion Models via Parameter-Efficient Post-Tuning
• Sample Code: Please refer to the older version
• Model: ModelScope, HuggingFace

https://github.com/modelscope/DiffSynth-Studio/assets/35051019/d97f6aa9-8064-4b5b-9d49-ed6001bb9acc

Diffutoon: High-Resolution Anime-Style Video Rendering

I don't care what you say, I just love anime!

• Project Page: Project Page
• Paper: Diffutoon: High-Resolution Editable Toon Shading via Diffusion Models
• Sample Code: Please refer to the older version

The Diffutoon examples are from the pre-2.0 codebase and are not included in the current main branch. Use the older source tree above with the corresponding older version of DiffSynth-Studio.

https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/b54c05c5-d747-4709-be5e-b39af82404dd

DiffSynth: The Original Version of This Project

In the era before video generation models, how could image generation models be used to process videos?

• Project Page: Project Page
• Paper: DiffSynth: Latent In-Iteration Deflickering for Realistic Video Synthesis
• Sample Code: Please refer to the older version

https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-4481-b79f-0c3a7361a1ea

Contact Us Copy Chapter

All Supported Models Copy Chapter

Overview (`docs/en/README.md`) Copy Section

DiffSynth-Studio Documentation

Welcome to the magical world of Diffusion models! DiffSynth-Studio is an open-source Diffusion model engine developed and maintained by the ModelScope Community. We aim to build a universal Diffusion model framework that fosters technological innovation through framework construction, aggregates the power of the open-source community, and explores the boundaries of generative model technology!

Documentation Reading Guide

graph LR;
    use["I want to use models for inference and training"]-->sec1["Section 1: Getting Started"];
    use["I want to use models for inference and training"]-->sec2["Section 2: Model Details"];
    use["I want to use models for inference and training"]-->sec3["Section 3: Training Framework"];
    develop["I want to develop based on this framework"]-->sec3["Section 3: Training Framework"];
    develop["I want to develop based on this framework"]-->sec4["Section 4: Model Integration"];
    develop["I want to develop based on this framework"]-->sec5["Section 5: API Reference"];
    explore["I want to explore new technologies based on this project"]-->sec4["Section 4: Model Integration"];
    explore["I want to explore new technologies based on this project"]-->sec5["Section 5: API Reference"];
    explore["I want to explore new technologies based on this project"]-->sec6["Section 6: Diffusion Templates"];
    explore["I want to explore new technologies based on this project"]-->sec7["Section 7: Research Guide"];
    problem["I encountered a problem"]-->sec8["Section 8: Frequently Asked Questions"];

Section 1: Getting Started Copy Chapter

This section introduces the basic usage of DiffSynth-Studio, including how to enable VRAM management for inference on GPUs with extremely low VRAM, and how to train various base models, LoRAs, ControlNets, and other models.

• Installation Dependencies
• Model Inference
• Accelerated Inference
• VRAM Management
• Model Quantization
• Model Training
• Environment Variables
• GPU/NPU Support
• Inference WebUI

Section 2: Model Details Copy Chapter

This section introduces the Diffusion models supported by DiffSynth-Studio. Some model pipelines feature special functionalities such as controllable generation and parallel acceleration.

• FLUX.1
• Wan
• Qwen-Image
• Qwen-Video-Edit
• FLUX.2
• Z-Image
• Anima
• LTX-2
• ERNIE-Image
• JoyAI-Image
• ACE-Step
• HiDream-O1-Image
• Stable Diffusion
• Stable Diffusion XL
• Image Quality Metrics
• Ideogram 4
• Krea-2
• Fibo 2.0
• Boogu-Image
• LingBot-Video
• MiniMax-H3
• MiniMax-Music3
• SenseNova-U1
• YuE2
• Qwen-Image-2.1
• FLUX.3

Section 3: Training Framework Copy Chapter

This section introduces the design philosophy of the training framework in DiffSynth-Studio, helping developers understand the principles of Diffusion model training algorithms.

• Basic Principles of Diffusion Models
• Standard Supervised Training
• Enabling FP8 Precision in Training
• End-to-End Distillation Accelerated Training
• Two-Stage Split Training
• Differential LoRA Training
• Enabling DeepSpeed
• Offload Training

Section 4: Model Integration Copy Chapter

This section introduces how to integrate models into DiffSynth-Studio to utilize the framework's basic functions, helping developers provide support for new models in this project or perform inference and training of private models.

• Integrating Model Architecture
• Building a Pipeline
• Enabling Fine-Grained VRAM Management
• Model Training Integration
• Integrating a Quantization Backend

We have open-sourced DiffSynth-Studio Model Integration Skills. This is a composable collection of Agent Skills that automates the entire workflow of integrating external diffusion models (image / video / audio) into DiffSynth-Studio. It distills and defines the model integration standards of DiffSynth-Studio, consolidating best practices such as codebase analysis, model code integration, Pipeline design, low-VRAM management, and training support into reusable standard procedures. Following these standards can significantly lower the integration barrier, reduce repetitive debugging, and greatly improve the efficiency of integrating new models. We recommend starting from the diffsynth-integrator example to get started quickly and accelerate model integration.

Section 5: API Reference Copy Chapter

This section introduces the independent core module diffsynth.core in DiffSynth-Studio, explaining how internal functions are designed and operate. Developers can use these functional modules in other codebase developments if needed.

• diffsynth.core.attention: Attention mechanism implementation
• diffsynth.core.data: Data processing operators and general datasets
• diffsynth.core.gradient: Gradient checkpointing
• diffsynth.core.loader: Model download and loading
• diffsynth.core.quant: Model quantization
• diffsynth.core.vram: VRAM management

Section 6: Diffusion Templates Copy Chapter

This section introduces the controllable generation plugin framework for Diffusion models, explaining the framework's operation mechanism and how to use Template models for inference and training.

• Introducing Diffusion Templates
• Diffusion Templates Architecture Details
• Template Model Inference
• Template Model Training

Section 7: Research Guide Copy Chapter

This section introduces how to use DiffSynth-Studio to train new models, helping researchers explore new model technologies.

• Training models from scratch
• Inference improvement techniques
• Designing controllable generation models
• Creating new training paradigms 【coming soon】

Section 8: Frequently Asked Questions Copy Chapter

This section summarizes common developer questions. If you encounter issues during usage or development, please refer to this section. If you still cannot resolve the problem, please submit an issue on GitHub.

• Frequently Asked Questions

Overview (`docs/zh/README.md`) Copy Section

DiffSynth-Studio 文档

欢迎来到 Diffusion 模型的魔法世界！DiffSynth-Studio 是由魔搭社区团队开发和维护的开源 Diffusion 模型引擎。我们期望构建一个通用的 Diffusion 模型框架，以框架建设孵化技术创新，凝聚开源社区的力量，探索生成式模型技术的边界！

文档阅读导引

graph LR;
    我想要使用模型进行推理和训练-->sec1[Section 1: 上手使用];
    我想要使用模型进行推理和训练-->sec2[Section 2: 模型详解];
    我想要使用模型进行推理和训练-->sec3[Section 3: 训练框架];
    我想要基于此框架进行二次开发-->sec3[Section 3: 训练框架];
    我想要基于此框架进行二次开发-->sec4[Section 4: 模型接入];
    我想要基于此框架进行二次开发-->sec5[Section 5: API 参考];
    我想要基于本项目探索新的技术-->sec4[Section 4: 模型接入];
    我想要基于本项目探索新的技术-->sec5[Section 5: API 参考];
    我想要基于本项目探索新的技术-->sec6[Section 6: Diffusion Templates]
    我想要基于本项目探索新的技术-->sec7[Section 7: 学术导引];
    我遇到了问题-->sec8[Section 8: 常见问题];

Section 1: 上手使用 Copy Chapter

本节介绍 DiffSynth-Studio 的基本使用方式，包括如何启用显存管理从而在极低显存的 GPU 上进行推理，以及如何训练任意基础模型、LoRA、ControlNet 等模型。

• 安装依赖
• 模型推理
• 加速推理
• 显存管理
• 模型量化
• 模型训练
• 环境变量
• GPU/NPU 支持
• 推理 WebUI

Section 2: 模型详解 Copy Chapter

本节介绍 DiffSynth-Studio 所支持的 Diffusion 模型，部分模型 Pipeline 具备可控生成、并行加速等特色功能。

• FLUX.1
• Wan
• Qwen-Image
• Qwen-Video-Edit
• FLUX.2
• Z-Image
• Anima
• LTX-2
• ERNIE-Image
• JoyAI-Image
• ACE-Step
• HiDream-O1-Image
• Stable Diffusion
• Stable Diffusion XL
• 图像质量评估指标
• Ideogram 4
• Krea-2
• Fibo 2.0
• Boogu-Image
• LingBot-Video
• MiniMax-H3
• MiniMax-Music3
• SenseNova-U1
• YuE2
• Qwen-Image-2.1
• FLUX.3

Section 3: 训练框架 Copy Chapter

本节介绍 DiffSynth-Studio 中训练框架的设计思路，帮助开发者理解 Diffusion 模型训练算法的原理。

• Diffusion 模型基本原理
• 标准监督训练
• 在训练中启用 FP8 精度
• 端到端的蒸馏加速训练
• 两阶段拆分训练
• 差分 LoRA 训练
• 启用 DeepSpeed
• Offload Training

Section 4: 模型接入 Copy Chapter

本节介绍如何将模型接入 DiffSynth-Studio 从而使用框架基础功能，帮助开发者为本项目提供新模型的支持，或进行私有化模型的推理和训练。

• 接入模型结构
• 接入 Pipeline
• 接入细粒度显存管理
• 接入模型训练
• 接入量化后端

我们开源了 DiffSynth-Studio Model Integration Skills。这是一套可组合的 Agent Skill 合集，将外部扩散模型（图像 / 视频 / 音频）接入 DiffSynth-Studio 的全流程自动化。它沉淀并定义了 DiffSynth-Studio 的模型接入规范，将代码库分析、模型代码接入、Pipeline 设计、低显存管理与训练支持等最佳实践固化为可复用的标准流程。遵循这套规范，能显著降低接入门槛、减少反复调试，大幅提升新模型的接入效率。建议从 diffsynth-integrator 的使用示例开始，快速上手、加速模型接入。

Section 5: API 参考 Copy Chapter

本节介绍 DiffSynth-Studio 中的独立核心模块 diffsynth.core，介绍内部的功能是如何设计和运作的，开发者如有需要，可将其中的功能模块用于其他代码库的开发中。

• diffsynth.core.attention: 注意力机制实现
• diffsynth.core.data: 数据处理算子与通用数据集
• diffsynth.core.gradient: 梯度检查点
• diffsynth.core.loader: 模型下载与加载
• diffsynth.core.quant: 模型量化
• diffsynth.core.vram: 显存管理

Section 6: Diffusion Templates Copy Chapter

本节介绍 Diffusion 模型可控生成插件框架 Diffusion Templates，讲解 Diffusion Templates 框架的运行机制，展示如何使用 Template 模型进行推理和训练。

• Diffusion Templates 简介
• Diffusion Templates 架构详解
• Template 模型推理
• Template 模型训练

Section 7: 学术导引 Copy Chapter

本节介绍如何利用 DiffSynth-Studio 训练新的模型，帮助科研工作者探索新的模型技术。

• 从零开始训练模型
• 推理改进优化技术
• 设计可控生成模型
• 创建新的训练范式【coming soon】

Section 8: 常见问题 Copy Chapter

本节总结了开发者常见的问题，如果你在使用和开发中遇到了问题，请参考本节内容，如果仍无法解决，请到 GitHub 上给我们提 issue。

• 常见问题

Overview (`examples/anima/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Anima.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Anima.html

Overview (`examples/fibo/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Fibo.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Fibo.html

Inference Copy Chapter

• Text-to-image, editing, masked editing, text rendering: model_inference/Fibo-2.0-dev.py
• Editing and masked editing: model_inference/Fibo-2.0-edit-turbo.py
• Low VRAM versions: model_inference_low_vram/

Training Copy Chapter

Only Fibo-2.0-dev is trainable.

Overview (`examples/flux/README.md`) Copy Section

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

Overview (`examples/flux2/README.md`) Copy Section

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

Overview (`examples/lingbot_video/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LingBot-Video.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/LingBot-Video.html

Overview (`examples/ltx2/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/LTX-2.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/LTX-2.html

Overview (`examples/mova/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Wan.html

Overview (`examples/qwen_image/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Qwen-Image.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Qwen-Image.html

Overview (`examples/wanvideo/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Wan.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Wan.html

Overview (`examples/yue2/README.md`) Copy Section

YuE2

YuE2 is a music generation model with editable symbolic planning.

• Inference: YuE2.py
• Low VRAM inference: YuE2.py
• Plan and render an edited score: YuE2-Edit.py
• Low VRAM plan and render an edited score: YuE2-Edit.py
• Documentation: English, Chinese

Overview (`examples/z_image/README.md`) Copy Section

English Document: https://diffsynth-studio-doc.readthedocs.io/en/latest/Model_Details/Z-Image.html

中文文档：https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Model_Details/Z-Image.html

Overview (`diffsynth/utils/ses/README.md`) Copy Section

Please see docs/en/Research_Tutorial/inference_time_scaling.md or docs/zh/Research_Tutorial/inference_time_scaling.md for more details.

Requirements (`docs/requirements.txt`) Copy Section

docutils>=0.16.0
myst_parser
recommonmark
sphinx>=5.3.0
sphinx-book-theme
sphinx-copybutton
sphinx-autobuild
sphinx-rtd-theme
sphinx_markdown_tables
sphinxcontrib-mermaid
pymdown-extensions

En .Readthedocs.Yaml (`docs/en/.readthedocs.yaml`) Copy Section

.readthedocs.yaml

Read the Docs configuration file

See https://docs.readthedocs.io/en/stable/config-file/v2.html for details

Required

version: 2

Set the OS, Python version and other tools you might need

build:
os: ubuntu-22.04
tools:
python: "3.10"

Build documentation in the "docs/" directory with Sphinx

sphinx:
configuration: docs/en/conf.py

Optionally build your docs in additional formats such as PDF and ePub

formats:

- pdf

- epub

Optional but recommended, declare the Python requirements required

to build your documentation

See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html

python:
install:
- requirements: docs/requirements.txt

En QA (`docs/en/QA.md`) Copy Section

Frequently Asked Questions

Why doesn't the training framework support batch size > 1? Copy Chapter

• Larger batch sizes no longer achieve significant acceleration: Due to acceleration technologies such as flash attention that have fully improved GPU utilization, larger batch sizes will only bring greater VRAM usage without significant acceleration. The experience with small models like Stable Diffusion 1.5 is no longer applicable to the latest large models.
• Larger batch sizes can be achieved through other solutions: Multi-GPU training and Gradient Accumulation can both mathematically equivalently achieve larger batch sizes.
• Larger batch sizes contradict the framework's general design: We hope to build a general training framework. Many models cannot accommodate larger batch sizes, such as text encodings of different lengths and images of different resolutions, which cannot be merged into larger batches.

Why aren't redundant parameters removed from certain models? Copy Chapter

In some models, redundant parameters exist. For example, in Qwen-Image's DiT model, the text portion of the last layer does not participate in any calculations. This is a minor bug left by the model developers. Setting it as trainable directly will also cause errors in multi-GPU training.

To maintain compatibility with other models in the open-source community, we have decided to retain these parameters. These redundant parameters can avoid errors in multi-GPU training through the --find_unused_parameters parameter.

Why does FP8 quantization show no acceleration effect? Copy Chapter

Native FP8 computation relies on Hopper architecture GPUs and has significant precision errors. It is currently immature technology, so this project does not support native FP8 computation.

FP8 computation in VRAM management refers to storing model parameters in memory or VRAM with FP8 precision and temporarily converting them to other precisions when needed for computation. Therefore, it can only reduce VRAM usage without acceleration effects.

Why doesn't the training framework support native FP8 precision training? Copy Chapter

Even with suitable hardware conditions, we currently have no plans to support native FP8 precision training.

• The main challenge of native FP8 precision training is precision overflow caused by gradient explosion. To ensure training stability, the model structure needs to be redesigned accordingly. However, no model developers are willing to do so at present.
• Additionally, models trained with native FP8 precision can only be computed with BF16 precision during inference without Hopper architecture GPUs, theoretically resulting in generation quality inferior to FP8.

Therefore, native FP8 precision training technology is extremely immature. We will observe the technological developments in the open-source community.

How to dynamically load LoRA models during inference? Copy Chapter

We support two loading methods for LoRA models. See LoRA Loading for details:

• Cold Loading: When VRAM Management is not enabled for the base model, LoRA will be fused into the base model weights. In this case, inference speed remains unchanged, and LoRA cannot be unloaded after loading.
• Hot Loading: When VRAM Management is enabled for the base model, LoRA will not be fused into the base model weights. In this case, inference speed will slow down, and LoRA can be unloaded after loading via pipe.clear_lora().

How to reduce VRAM usage during training? Copy Chapter

The framework supports multiple methods to reduce VRAM usage during training, including Gradient Checkpointing, DeepSpeed, FP8, Two-Stage Split Training, and CPU Offload. Please refer to Low VRAM Training.

En Index (`docs/en/index.rst`) Copy Section

Welcome to DiffSynth-Studio's Documentation

.. toctree::
:maxdepth: 2
:caption: Documentation Introduction

README

.. toctree::
:maxdepth: 2
:caption: Getting Started

Pipeline_Usage/Setup
Pipeline_Usage/Model_Inference
Pipeline_Usage/Accelerated_Inference
Pipeline_Usage/VRAM_management
Pipeline_Usage/Quantization
Pipeline_Usage/Model_Training
Pipeline_Usage/Environment_Variables
Pipeline_Usage/GPU_support
Pipeline_Usage/Inference_WebUI

.. toctree::
:maxdepth: 2
:caption: Model Details

Model_Details/FLUX
Model_Details/Wan
Model_Details/Qwen-Image
Model_Details/Qwen-Video-Edit
Model_Details/FLUX2
Model_Details/Z-Image
Model_Details/Anima
Model_Details/LTX-2
Model_Details/LTX-2.5
Model_Details/ERNIE-Image
Model_Details/JoyAI-Image
Model_Details/ACE-Step
Model_Details/HiDream-O1-Image
Model_Details/Stable-Diffusion
Model_Details/Stable-Diffusion-XL
Model_Details/Image-Quality-Metrics
Model_Details/Ideogram-4
Model_Details/Krea-2
Model_Details/Boogu-Image
Model_Details/LingBot-Video
Model_Details/MiniMax-H3
Model_Details/DiffSynth-Music
Model_Details/MiniMax-Music3
Model_Details/YuE2
Model_Details/SenseNova-U1
Model_Details/Qwen-Image-2.1
Model_Details/FLUX.3
Model_Details/MOSS-Image
Model_Details/Fibo

.. toctree::
:maxdepth: 2
:caption: Training Framework

Training/Understanding_Diffusion_models
Training/Supervised_Fine_Tuning
Training/FP8_Precision
Training/Distillation
Training/Direct_Distill
Training/Split_Training
Training/Differential_LoRA
Training/DeepSpeed
Training/Offload_Training

.. toctree::
:maxdepth: 2
:caption: Model Integration

Developer_Guide/Integrating_Your_Model
Developer_Guide/Building_a_Pipeline
Developer_Guide/Enabling_VRAM_management
Developer_Guide/Training_Diffusion_Models
Developer_Guide/Integrating_Quantization_Backend

.. toctree::
:maxdepth: 2
:caption: API Reference

API_Reference/core/attention
API_Reference/core/data
API_Reference/core/gradient
API_Reference/core/loader
API_Reference/core/quant
API_Reference/core/vram

.. toctree::
:maxdepth: 2
:caption: Diffusion Templates

Diffusion_Templates/Introducing_Diffusion_Templates.md
Diffusion_Templates/Understanding_Diffusion_Templates.md
Diffusion_Templates/Template_Model_Inference.md
Diffusion_Templates/Template_Model_Training.md

.. toctree::
:maxdepth: 2
:caption: Research Guide

Research_Tutorial/train_from_scratch
Research_Tutorial/inference_time_scaling
Research_Tutorial/controllable_models

.. toctree::
:maxdepth: 2
:caption: FAQ

QA

Indices and tables

• :ref:genindex
• :ref:modindex
• :ref:search

Zh .Readthedocs.Yaml (`docs/zh/.readthedocs.yaml`) Copy Section

.readthedocs.yaml

Read the Docs configuration file

See https://docs.readthedocs.io/en/stable/config-file/v2.html for details

Required

version: 2

Set the OS, Python version and other tools you might need

build:
os: ubuntu-22.04
tools:
python: "3.10"

Build documentation in the "docs/" directory with Sphinx

sphinx:
configuration: docs/zh/conf.py

Optionally build your docs in additional formats such as PDF and ePub

formats:

- pdf

- epub

Optional but recommended, declare the Python requirements required

to build your documentation

See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html

python:
install:
- requirements: docs/requirements.txt

Zh QA (`docs/zh/QA.md`) Copy Section

常见问题

为什么训练框架不支持 batch size > 1？ Copy Chapter

• 更大的 batch size 已无法实现显著加速：由于 flash attention 等加速技术已经充分提高了 GPU 的利用率，因此更大的 batch size 只会带来更大的显存占用，无法带来显著加速。在 Stable Diffusion 1.5 这类小模型上的经验已不再适用于最新的大模型。
• 更大的 batch size 可以用其他方案实现：多 GPU 训练和 Gradient Accumulation 都可以在数学意义上等价地实现更大的 batch size。
• 更大的 batch size 与框架的通用性设计相悖：我们希望构建通用的训练框架，大量模型无法适配更大的 batch size，例如不同长度的文本编码、不同分辨率的图像等，都是无法合并为更大的 batch 的。

为什么不删除某些模型中的冗余参数？ Copy Chapter

在部分模型中，模型存在冗余参数，例如 Qwen-Image 的 DiT 模型最后一层的文本部分，这部分参数不会参与任何计算，这是模型开发者留下的小 bug。直接将其设置为可训练时还会在多 GPU 训练中出现报错。

为了与开源社区中其他模型保持兼容性，我们决定保留这些参数。这些冗余参数在多 GPU 训练中可以通过 --find_unused_parameters 参数避免报错。

为什么 FP8 量化没有任何加速效果？ Copy Chapter

原生 FP8 计算需要依赖 Hopper 架构的 GPU，同时在计算精度上有较大误差，目前仍然是不成熟的技术，因此本项目不支持原生 FP8 计算。

显存管理中的 FP8 计算是指将模型参数以 FP8 精度存储在内存或显存中，在需要计算时临时转换为其他精度，因此仅能减少显存占用，没有加速效果。

为什么训练框架不支持原生 FP8 精度训练？ Copy Chapter

即使硬件条件允许，我们目前也没有任何支持原生 FP8 精度训练的规划。

• 目前原生 FP8 精度训练的主要挑战是梯度爆炸导致的精度溢出，为了保证训练的稳定性，需针对性地重新设计模型结构，然而目前还没有任何模型开发者愿意这么做。
• 此外，使用原生 FP8 精度训练的模型，在推理时若没有 Hopper 架构 GPU，则只能以 BF16 精度进行计算，理论上其生成效果反而不如 FP8。

因此，原生 FP8 精度训练技术是极不成熟的，我们静观开源社区的技术发展。

如何在推理时动态加载 LoRA 模型？ Copy Chapter

我们支持 LoRA 模型的两种加载方式，详见LoRA 加载：

• 冷加载：当基础模型未开启显存管理时，LoRA 会融合进基础模型权重，此时推理速度没有变化，LoRA 加载后无法卸载。
• 热加载：当基础模型开启显存管理时，LoRA 不会融合进基础模型权重，此时推理速度会变慢，LoRA 加载后可通过 pipe.clear_lora() 卸载。

如何减少训练所需的显存？ Copy Chapter

框架支持多种方式减少训练所需的显存，包括 Gradient Checkpointing、DeepSpeed、FP8、两阶段拆分训练、CPU Offload，请参考低显存训练。

Zh Index (`docs/zh/index.rst`) Copy Section

欢迎来到 DiffSynth-Studio 的文档

.. toctree::
:maxdepth: 2
:caption: 文档介绍

README

.. toctree::
:maxdepth: 2
:caption: 上手使用

Pipeline_Usage/Setup
Pipeline_Usage/Model_Inference
Pipeline_Usage/Accelerated_Inference
Pipeline_Usage/VRAM_management
Pipeline_Usage/Quantization
Pipeline_Usage/Model_Training
Pipeline_Usage/Environment_Variables
Pipeline_Usage/GPU_support
Pipeline_Usage/Inference_WebUI

.. toctree::
:maxdepth: 2
:caption: 模型详解

Model_Details/FLUX
Model_Details/Wan
Model_Details/Qwen-Image
Model_Details/Qwen-Video-Edit
Model_Details/FLUX2
Model_Details/Z-Image
Model_Details/Anima
Model_Details/LTX-2
Model_Details/LTX-2.5
Model_Details/ERNIE-Image
Model_Details/JoyAI-Image
Model_Details/ACE-Step
Model_Details/HiDream-O1-Image
Model_Details/Stable-Diffusion
Model_Details/Stable-Diffusion-XL
Model_Details/Image-Quality-Metrics
Model_Details/Ideogram-4
Model_Details/Krea-2
Model_Details/Boogu-Image
Model_Details/LingBot-Video
Model_Details/MiniMax-H3
Model_Details/DiffSynth-Music
Model_Details/MiniMax-Music3
Model_Details/YuE2
Model_Details/SenseNova-U1
Model_Details/Qwen-Image-2.1
Model_Details/FLUX.3
Model_Details/MOSS-Image
Model_Details/Fibo

.. toctree::
:maxdepth: 2
:caption: 训练框架

Training/Understanding_Diffusion_models
Training/Supervised_Fine_Tuning
Training/FP8_Precision
Training/Distillation
Training/Direct_Distill
Training/Split_Training
Training/Differential_LoRA
Training/DeepSpeed
Training/Offload_Training

.. toctree::
:maxdepth: 2
:caption: 模型接入

Developer_Guide/Integrating_Your_Model
Developer_Guide/Building_a_Pipeline
Developer_Guide/Enabling_VRAM_management
Developer_Guide/Training_Diffusion_Models
Developer_Guide/Integrating_Quantization_Backend

.. toctree::
:maxdepth: 2
:caption: API 参考

API_Reference/core/attention
API_Reference/core/data
API_Reference/core/gradient
API_Reference/core/loader
API_Reference/core/quant
API_Reference/core/vram

.. toctree::
:maxdepth: 2
:caption: Diffusion Templates

Diffusion_Templates/Introducing_Diffusion_Templates.md
Diffusion_Templates/Understanding_Diffusion_Templates.md
Diffusion_Templates/Template_Model_Inference.md
Diffusion_Templates/Template_Model_Training.md

.. toctree::
:maxdepth: 2
:caption: 学术导引

Research_Tutorial/train_from_scratch
Research_Tutorial/inference_time_scaling
Research_Tutorial/controllable_models

.. toctree::
:maxdepth: 2
:caption: 常见问题

QA

Indices and tables

• :ref:genindex
• :ref:modindex
• :ref:search

En Developer Guide Building A Pipeline (`docs/en/Developer_Guide/Building_a_Pipeline.md`) Copy Section

Building a Pipeline

After integrating the required models for the Pipeline, you also need to build a Pipeline for model inference. This document provides a standardized process for building a Pipeline. Developers can also refer to existing Pipeline implementations for construction.

The Pipeline implementation is located in diffsynth/pipelines. Each Pipeline contains the following essential key components:

• __init__
• from_pretrained
• __call__
• units
• model_fn

`__init__` Copy Chapter

In __init__, the Pipeline is initialized. Here is a simple implementation:

import torch
from PIL import Image
from typing import Union
from tqdm import tqdm
from ..diffusion import FlowMatchScheduler
from ..core import ModelConfig
from ..diffusion.base_pipeline import BasePipeline, PipelineUnit
from ..models.new_models import XXX_Model, YYY_Model, ZZZ_Model

class NewDiffSynthPipeline(BasePipeline):

    def __init__(self, device="cuda", torch_dtype=torch.bfloat16):
        super().__init__(device=device, torch_dtype=torch_dtype)
        self.scheduler = FlowMatchScheduler()
        self.text_encoder: XXX_Model = None
        self.dit: YYY_Model = None
        self.vae: ZZZ_Model = None
        self.in_iteration_models = ("dit",)
        self.units = [
            NewDiffSynthPipelineUnit_xxx(),
            ...
        ]
        self.model_fn = model_fn_new

This includes the following parts:

• scheduler: Scheduler, used to control the coefficients in the iterative formula during inference, controlling the noise content at each step.
• text_encoder, dit, vae: Models. Since Latent Diffusion was proposed, this three-stage model architecture has become the mainstream Diffusion model architecture. However, this is not immutable, and any number of models can be added to the Pipeline.
• in_iteration_models: Iteration models. This tuple marks which models will be called during iteration.
• units: Pre-processing units for model iteration. See units for details.
• model_fn: The forward function of the denoising model during iteration. See model_fn for details.

Q: Model loading does not occur in __init__, why initialize each model as None here?

A: By annotating the type of each model here, the code editor can provide code completion prompts based on each model, facilitating subsequent development.

`from_pretrained` Copy Chapter

from_pretrained is responsible for loading the required models to make the Pipeline callable. Here is a simple implementation:

@staticmethod
    def from_pretrained(
        torch_dtype: torch.dtype = torch.bfloat16,
        device: Union[str, torch.device] = "cuda",
        model_configs: list[ModelConfig] = [],
        vram_limit: float = None,
    ):
        # Initialize pipeline
        pipe = NewDiffSynthPipeline(device=device, torch_dtype=torch_dtype)
        model_pool = pipe.download_and_load_models(model_configs, vram_limit)
        
        # Fetch models
        pipe.text_encoder = model_pool.fetch_model("xxx_text_encoder")
        pipe.dit = model_pool.fetch_model("yyy_dit")
        pipe.vae = model_pool.fetch_model("zzz_vae")
        # If necessary, load tokenizers here.
        
        # VRAM Management
        pipe.vram_management_enabled = pipe.check_vram_management_state()
        return pipe

Developers need to implement the logic for fetching models. The corresponding model names are the "model_name" in the model Config filled in during model integration.

Some models also need to load tokenizer. Extra tokenizer_config parameters can be added to from_pretrained as needed, and this part can be implemented after fetching the models.

`__call__` Copy Chapter

__call__ implements the entire generation process of the Pipeline. Below is a common generation process template. Developers can modify it based on their needs.

/* Detailed source-code truncated for AI context efficiency. */

`units` Copy Chapter

units contains all the preprocessing processes, such as: width/height checking, prompt encoding, initial noise generation, etc. In the entire model preprocessing process, data is abstracted into three mutually exclusive parts, stored in corresponding dictionaries:

• inputs_shared: Shared inputs, parameters unrelated to Classifier-Free Guidance (CFG for short).
• inputs_posi: Positive side inputs for Classifier-Free Guidance, containing content related to positive prompts.
• inputs_nega: Negative side inputs for Classifier-Free Guidance, containing content related to negative prompts.

Pipeline Unit implementations include three types: direct mode, CFG separation mode, and takeover mode.

If some calculations are unrelated to CFG, direct mode can be used, for example, Qwen-Image's random noise initialization:

class QwenImageUnit_NoiseInitializer(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width", "seed", "rand_device"),
            output_params=("noise",),
        )

    def process(self, pipe: QwenImagePipeline, height, width, seed, rand_device):
        noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype)
        return {"noise": noise}

If some calculations are related to CFG and need to separately process positive and negative prompts, but the input parameters on both sides are the same, CFG separation mode can be used, for example, Qwen-image's prompt encoding:

class QwenImageUnit_PromptEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            seperate_cfg=True,
            input_params_posi={"prompt": "prompt"},
            input_params_nega={"prompt": "negative_prompt"},
            input_params=("edit_image",),
            output_params=("prompt_emb", "prompt_emb_mask"),
            onload_model_names=("text_encoder",)
        )

    def process(self, pipe: QwenImagePipeline, prompt, edit_image=None) -> dict:
        pipe.load_models_to_device(self.onload_model_names)
        # Do something
        return {"prompt_emb": prompt_embeds, "prompt_emb_mask": encoder_attention_mask}

If some calculations need global information, takeover mode is required, for example, Qwen-Image's entity partition control:

class QwenImageUnit_EntityControl(PipelineUnit):
    def __init__(self):
        super().__init__(
            take_over=True,
            input_params=("eligen_entity_prompts", "width", "height", "eligen_enable_on_negative", "cfg_scale"),
            output_params=("entity_prompt_emb", "entity_masks", "entity_prompt_emb_mask"),
            onload_model_names=("text_encoder",)
        )

    def process(self, pipe: QwenImagePipeline, inputs_shared, inputs_posi, inputs_nega):
        # Do something
        return inputs_shared, inputs_posi, inputs_nega

The following are the parameter configurations required for Pipeline Unit:

• seperate_cfg: Whether to enable CFG separation mode
• take_over: Whether to enable takeover mode
• input_params: Shared input parameters
• output_params: Output parameters
• input_params_posi: Positive side input parameters
• input_params_nega: Negative side input parameters
• onload_model_names: Names of model components to be called

When designing unit, please try to follow these principles:

• Default fallback: For optional function unit input parameters, the default is None rather than False or other values. Please provide fallback processing for this default value.
• Parameter triggering: Some Adapter models may not be loaded, such as ControlNet. The corresponding unit should control triggering based on whether the parameter input is None rather than whether the model is loaded. For example, when the user inputs controlnet_image but does not load the ControlNet model, the code should give an error rather than ignore these input parameters and continue execution.
• Simplicity first: Use direct mode as much as possible, only use takeover mode when the function cannot be implemented.
• VRAM efficiency: When calling models in unit, please use pipe.load_models_to_device(self.onload_model_names) to activate the corresponding models. Do not call other models outside onload_model_names. After unit calculation is completed, do not manually release VRAM with pipe.load_models_to_device([]).

Q: Some parameters are not called during the inference process, such as output_params. Is it still necessary to configure them?

A: These parameters will not affect the inference process, but they will affect some experimental features. Therefore, we recommend configuring them properly. For example, "split training" - we can complete the preprocessing offline during training, but some model calculations that require gradient backpropagation cannot be split. These parameters are used to build computational graphs to infer which calculations can be split.

`model_fn` Copy Chapter

model_fn is the unified forward interface during iteration. For models where the open-source ecosystem is not yet formed, you can directly use the denoising model's forward, for example:

def model_fn_new(dit=None, latents=None, timestep=None, prompt_emb=None, **kwargs):
    return dit(latents, prompt_emb, timestep)

For models with rich open-source ecosystems, model_fn usually contains complex and chaotic cross-model inference. Taking diffsynth/pipelines/qwen_image.py as an example, the additional calculations implemented in this function include: entity partition control, three types of ControlNet, Gradient Checkpointing, etc. Developers need to be extra careful when implementing this part to avoid conflicts between module functions.

Compilation Acceleration Copy Chapter

To enable compilation acceleration, please refer to Inference Acceleration.

En Developer Guide Enabling VRAM Management (`docs/en/Developer_Guide/Enabling_VRAM_management.md`) Copy Section

Fine-Grained VRAM Management Scheme

This document introduces how to write reasonable fine-grained VRAM management schemes for models, and how to use the VRAM management functions in DiffSynth-Studio for other external code libraries. Before reading this document, please read the document VRAM Management.

How Much VRAM Does a 20B Model Need? Copy Chapter

Taking Qwen-Image's DiT model as an example, this model has reached 20B parameters. The following code will load this model and perform inference, requiring about 40G VRAM. This model obviously cannot run on consumer-grade GPUs with smaller VRAM.

from diffsynth.core import load_model
from diffsynth.models.qwen_image_dit import QwenImageDiT
from modelscope import snapshot_download
import torch

snapshot_download(
    model_id="Qwen/Qwen-Image",
    local_dir="models/Qwen/Qwen-Image",
    allow_file_pattern="transformer/*"
)
prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(QwenImageDiT, model_path, torch_dtype=torch.bfloat16, device="cuda")
with torch.no_grad():
    output = model(**inputs)

Writing Fine-Grained VRAM Management Scheme Copy Chapter

To write a fine-grained VRAM management scheme, we need to use print(model) to observe and analyze the model structure:

/* Detailed source-code truncated for AI context efficiency. */

In VRAM management, we only care about layers containing parameters. In this model structure, QwenEmbedRope, TemporalTimesteps, SiLU and other Layers do not contain parameters. LayerNorm also does not contain parameters because elementwise_affine=False is set. Layers containing parameters are only Linear and RMSNorm.

diffsynth.core.vram provides two replacement modules for VRAM management:

• AutoWrappedLinear: Used to replace Linear layers
• AutoWrappedModule: Used to replace any other layer

Write a module_map to map Linear and RMSNorm in the model to the corresponding modules:

module_map={
    torch.nn.Linear: AutoWrappedLinear,
    RMSNorm: AutoWrappedModule,
}

In addition, vram_config and vram_limit are also required, which have been introduced in VRAM Management.

Call enable_vram_management to enable VRAM management. Note that the device when loading the model is cpu, consistent with offload_device:

from diffsynth.core import load_model, enable_vram_management, AutoWrappedLinear, AutoWrappedModule
from diffsynth.models.qwen_image_dit import QwenImageDiT, RMSNorm
import torch

prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(QwenImageDiT, model_path, torch_dtype=torch.bfloat16, device="cpu")
enable_vram_management(
    model,
    module_map={
        torch.nn.Linear: AutoWrappedLinear,
        RMSNorm: AutoWrappedModule,
    },
    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",
    },
    vram_limit=0,
)
with torch.no_grad():
    output = model(**inputs)

The above code only requires 2G VRAM to run the forward of a 20B model.

Disk Offload Copy Chapter

Disk Offload is a special VRAM management scheme that needs to be enabled during the model loading process, not after the model is loaded. Usually, when the above code can run smoothly, Disk Offload can be directly enabled:

from diffsynth.core import load_model, enable_vram_management, AutoWrappedLinear, AutoWrappedModule
from diffsynth.models.qwen_image_dit import QwenImageDiT, RMSNorm
import torch

prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(
    QwenImageDiT,
    model_path,
    module_map={
        torch.nn.Linear: AutoWrappedLinear,
        RMSNorm: AutoWrappedModule,
    },
    vram_config={
        "offload_dtype": "disk",
        "offload_device": "disk",
        "onload_dtype": "disk",
        "onload_device": "disk",
        "preparing_dtype": torch.bfloat16,
        "preparing_device": "cuda",
        "computation_dtype": torch.bfloat16,
        "computation_device": "cuda",
    },
    vram_limit=0,
)
with torch.no_grad():
    output = model(**inputs)

Disk Offload is an extremely special VRAM management scheme. It only supports .safetensors format files, not binary files such as .bin, .pth, .ckpt, and does not support state dict converter with Tensor reshape.

If there are situations where Disk Offload cannot run normally but non-Disk Offload can run normally, please submit an issue to us on GitHub.

Writing Default Configuration Copy Chapter

To make it easier for users to use the VRAM management function, we write the fine-grained VRAM management configuration in diffsynth/configs/vram_management_module_maps.py. The configuration information for the above model is:

"diffsynth.models.qwen_image_dit.QwenImageDiT": {
    "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
    "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
}
```# Fine-Grained VRAM Management Scheme

This document introduces how to write reasonable fine-grained VRAM management schemes for models, and how to use the VRAM management functions in `DiffSynth-Studio` for other external code libraries. Before reading this document, please read the document [VRAM Management](../Pipeline_Usage/VRAM_management.md).

## How Much VRAM Does a 20B Model Need?

Taking Qwen-Image's DiT model as an example, this model has reached 20B parameters. The following code will load this model and perform inference, requiring about 40G VRAM. This model obviously cannot run on consumer-grade GPUs with smaller VRAM.

```python
from diffsynth.core import load_model
from diffsynth.models.qwen_image_dit import QwenImageDiT
from modelscope import snapshot_download
import torch

snapshot_download(
    model_id="Qwen/Qwen-Image",
    local_dir="models/Qwen/Qwen-Image",
    allow_file_pattern="transformer/*"
)
prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(QwenImageDiT, model_path, torch_dtype=torch.bfloat16, device="cuda")
with torch.no_grad():
    output = model(**inputs)

Writing Fine-Grained VRAM Management Scheme Copy Chapter

To write a fine-grained VRAM management scheme, we need to use print(model) to observe and analyze the model structure:

/* Detailed source-code truncated for AI context efficiency. */

In VRAM management, we only care about layers containing parameters. In this model structure, QwenEmbedRope, TemporalTimesteps, SiLU and other Layers do not contain parameters. LayerNorm also does not contain parameters because elementwise_affine=False is set. Layers containing parameters are only Linear and RMSNorm.

diffsynth.core.vram provides two replacement modules for VRAM management:

• AutoWrappedLinear: Used to replace Linear layers
• AutoWrappedModule: Used to replace any other layer

Write a module_map to map Linear and RMSNorm in the model to the corresponding modules:

module_map={
    torch.nn.Linear: AutoWrappedLinear,
    RMSNorm: AutoWrappedModule,
}

In addition, vram_config and vram_limit are also required, which have been introduced in VRAM Management.

Call enable_vram_management to enable VRAM management. Note that the device when loading the model is cpu, consistent with offload_device:

from diffsynth.core import load_model, enable_vram_management, AutoWrappedLinear, AutoWrappedModule
from diffsynth.models.qwen_image_dit import QwenImageDiT, RMSNorm
import torch

prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(QwenImageDiT, model_path, torch_dtype=torch.bfloat16, device="cpu")
enable_vram_management(
    model,
    module_map={
        torch.nn.Linear: AutoWrappedLinear,
        RMSNorm: AutoWrappedModule,
    },
    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",
    },
    vram_limit=0,
)
with torch.no_grad():
    output = model(**inputs)

The above code only requires 2G VRAM to run the forward of a 20B model.

Disk Offload Copy Chapter

Disk Offload is a special VRAM management scheme that needs to be enabled during the model loading process, not after the model is loaded. Usually, when the above code can run smoothly, Disk Offload can be directly enabled:

from diffsynth.core import load_model, enable_vram_management, AutoWrappedLinear, AutoWrappedModule
from diffsynth.models.qwen_image_dit import QwenImageDiT, RMSNorm
import torch

prefix = "models/Qwen/Qwen-Image/transformer/diffusion_pytorch_model"
model_path = [prefix + f"-0000{i}-of-00009.safetensors" for i in range(1, 10)]
inputs = {
    "latents": torch.randn((1, 16, 128, 128), dtype=torch.bfloat16, device="cuda"),
    "timestep": torch.zeros((1,), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb": torch.randn((1, 5, 3584), dtype=torch.bfloat16, device="cuda"),
    "prompt_emb_mask": torch.ones((1, 5), dtype=torch.int64, device="cuda"),
    "height": 1024,
    "width": 1024,
}

model = load_model(
    QwenImageDiT,
    model_path,
    module_map={
        torch.nn.Linear: AutoWrappedLinear,
        RMSNorm: AutoWrappedModule,
    },
    vram_config={
        "offload_dtype": "disk",
        "offload_device": "disk",
        "onload_dtype": "disk",
        "onload_device": "disk",
        "preparing_dtype": torch.bfloat16,
        "preparing_device": "cuda",
        "computation_dtype": torch.bfloat16,
        "computation_device": "cuda",
    },
    vram_limit=0,
)
with torch.no_grad():
    output = model(**inputs)

Disk Offload is an extremely special VRAM management scheme. It only supports .safetensors format files, not binary files such as .bin, .pth, .ckpt, and does not support state dict converter with Tensor reshape.

If there are situations where Disk Offload cannot run normally but non-Disk Offload can run normally, please submit an issue to us on GitHub.

Writing Default Configuration Copy Chapter

To make it easier for users to use the VRAM management function, we write the fine-grained VRAM management configuration in diffsynth/configs/vram_management_module_maps.py. The configuration information for the above model is:

"diffsynth.models.qwen_image_dit.QwenImageDiT": {
    "diffsynth.models.qwen_image_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
    "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
}

En Developer Guide Integrating Quantization Backend (`docs/en/Developer_Guide/Integrating_Quantization_Backend.md`) Copy Section

Integrating a Quantization Backend

The quantization framework of DiffSynth-Studio lives in diffsynth.core.quant and ships with bitsandbytes, torchao, and comfy-kitchen backends (see Model Quantization). If you have your own quantization algorithm, or want to plug in another quantization library, you only need to implement a QuantBackend — online quantization, saving/loading pre-quantized checkpoints, mixed quantization, VRAM management, and quantization + LoRA training are all reused as-is.

This guide walks through the whole process with a toy backend: INT9 — 9-bit symmetric weight-only quantization, genuinely stored at 9 bits per weight, with one fp32 scale per output channel. INT9 does not exist on any hardware; it is used here because it keeps the example short while still covering every interface you have to implement. For the full interface signatures and contracts, see the diffsynth.core.quant API documentation.

Framework Structure Copy Chapter

The framework has three layers:

• QuantizeConfig: the user-facing config and entry point, responsible for traversing the model, matching layers, and replacing nn.Linear. You never need to touch it.
• QuantBackend: the adapter layer, which only ever deals with a single nn.Linear: how to quantize it, how to build an empty shell, how to dequantize it, how to read and write its state dict. This is the part you implement.
• The quantized Linear: the module that actually holds the quantized weight and performs dequantization + matmul in forward.

The quantized Linear must satisfy four contract clauses:

• (a) It is a drop-in replacement for nn.Linear, with forward(x) doing dequantization + matmul internally. It must subclass torch.nn.Linear, otherwise LoRA injection and VRAM management cannot see it.
• (b) .to(...) only moves devices, never re-types the packed weight or quant state. VRAM management performs dtype conversions on the model; if a packed weight were cast to bf16, the quant state would be corrupted.
• (c) state_dict() and load_state_dict(assign=True) round-trip, via flatten_state_dict / unflatten_state_dict when necessary.
• (d) (Training only) forward is differentiable w.r.t. its input, so gradients can pass through frozen quantized layers to reach LoRA branches.

Step 1: Write the Quantized Linear Copy Chapter

INT9's storage layout needs a little thought: there is no native 9-bit dtype, and simply putting the codes into an int16 tensor would still spend 16 bits per weight — exactly as much as bf16, so the quantization would save nothing. Each weight is therefore split in two: the low 8 bits go into a uint8 weight buffer, and the 9th (most significant) bit forms a separate bit plane where 8 weights are packed into one byte in weight_msb, plus one fp32 weight_scale per output channel. That is 9 bits per weight, 56% of bf16.

Two more details matter:

• Delete nn.Linear's original weight parameter and register a buffer under the same name, so checkpoint keys stay layer_name.weight. Disk offload and mixed-quantization key ownership rely on this.
• Guard the dtype of the packed tensors by overriding _apply, i.e. contract clause (b). Every conversion (.to(), .half(), .float(), ...) funnels through _apply, so a conversion that would change the dtype is downgraded to a device-only move.

/* Detailed source-code truncated for AI context efficiency. */

Dequantization in forward uses ordinary tensor ops, so gradients flow back to the input x through F.linear and contract clause (d) holds automatically. Unpacking here is written bit by bit with PyTorch ops purely for clarity; a real backend fuses unpacking into the matmul kernel instead of materializing an fp weight on every forward.

There is an easy trap in dequantize_weight: the integer codes must be reconstructed in fp32. bf16 only carries 8 bits of significand, so integers above 256 are not representable exactly; casting the codes to bf16 before applying the scale rounds the 9th bit away and throws the accuracy gain out (measured: the advantage over int8 collapses from 2.25x to 1.15x). This applies to any format whose code width exceeds the significand of the compute dtype.

Step 2: Write the Backend Copy Chapter

Every backend method operates on a single nn.Linear:

• capabilities(): declares what the backend supports; all four flags default to False. Saving quantized weights requires is_serializable=True, and quantization + LoRA training requires is_differentiable=True.
• quantized_linear_classes(): declares the Linear classes this backend produces; is_quantized_linear defaults to an isinstance check against them.
• create_quantized_linear(): online quantization, turning an fp nn.Linear into a quantized one. compute_device is where quantization runs and model_device is where the result is stored, so the two together stream the work layer by layer with only one layer on the accelerator at a time.
• create_quantized_linear_shell(): builds an empty shell, used for loading pre-quantized checkpoints and for disk offload. It is rebuilt on every offload cycle, so build it on the meta device and keep it cheap.
• dequantize_to_linear(): restores a plain nn.Linear, used by mode="dequant_once".
• flatten_state_dict / unflatten_state_dict: conversion between the state dict and flat tensors. INT9's state dict already holds plain tensors, so the base class implementation is enough; only backends with composite tensors (tensor subclasses, nested quant state) such as bitsandbytes and torchao need to override them.

Unimplemented methods raise a descriptive exception from the base class, so a backend that only supports some capabilities can implement just what it needs. self.config is the backend config instance injected by the framework — the Int9WeightOnlyConfig written in the next step.

/* Detailed source-code truncated for AI context efficiency. */

Step 3: Write the Backend Config Copy Chapter

The backend config subclasses BackendConfig: user-tunable parameters are ordinary dataclass fields, while values pinned by the method are declared with field(init=False, default=...). describe_quant_method reports the two groups separately, and from_kwargs raises when a user passes unknown backend_config_kwargs.

@dataclass
class Int9WeightOnlyConfig(BackendConfig):
    per_channel: bool = True                     # user-tunable: per-channel or per-tensor
    bits: int = field(init=False, default=9)      # pinned by the method, not overridable

Step 4: Register the Quantization Method Copy Chapter

One backend can register several methods, distinguished by the fields pinned in its config (the bitsandbytes backend, for example, distinguishes nf4 from fp4 via quant_type). Method names should follow the _<format>_w<weight bits>a<activation bits> convention:

register_quant_method("toy_int9_w9a16", "toy_int9", Int9WeightOnlyConfig.from_kwargs, label="9bit, int9, weight-only (toy)")

There are two ways to register a backend and its methods:

Option 1: keep it in your own code (recommended, plug-and-play). Put the code above in any module; as long as that module is imported before you construct QuantizeConfig, the method is already in QUANT_METHODS and can be used just like a built-in one, with no framework changes:

import my_project.toy_int9   # triggers register_quant_backend / register_quant_method

from diffsynth.core.quant import QuantizeConfig

quantize = QuantizeConfig(method="toy_int9_w9a16", backend_config_kwargs={"per_channel": True})

Option 2: ship it as a built-in backend (permanent). Put the backend file under diffsynth/core/quant/backends/ and register it in _LAZY_BACKENDS in diffsynth/core/quant/backends/__init__.py; the framework then imports it on demand and users do not need to import anything:

_LAZY_BACKENDS = {
    "bitsandbytes": ".bitsandbytes",
    "torchao": ".torchao",
    "comfy_kitchen": ".comfy_kitchen",
    "toy_int9": ".toy_int9",
}

If your quantization algorithm or library is generally useful, you are welcome to submit it as a PR following Option 2, so that more users can benefit from it. A backend that depends on a third-party library should check its dependencies in validate_environment() with an installation hint, and point project_url at the upstream project.

Step 5: Self-Check Copy Chapter

The framework provides two verification tools; run them right after integrating. check_backend_contract verifies that the backend declares its Linear classes, that both factory methods return instances of those classes, that all declared classes subclass nn.Linear, and that every checkpoint key the backend actually writes lives under the layer name (a missing scale would make disk offload silently load corrupted layers). Unsupported factory methods are skipped rather than counted as failures.

from diffsynth.core.quant import QUANT_BACKENDS, QUANT_METHODS, check_backend_contract, check_differentiable, describe_quant_method

describe_quant_method("toy_int9_w9a16")

spec = QUANT_METHODS["toy_int9_w9a16"]
check_backend_contract(QUANT_BACKENDS[spec.backend](spec.config_factory({})), compute_device="cpu")

The output is as follows; describe_quant_method also confirms that the split between user-tunable and pinned parameters is what you intended:

method: toy_int9_w9a16
backend: toy_int9
detail: 9bit, int9, weight-only (toy)
backend config: my_project.toy_int9.Int9WeightOnlyConfig
backend_config_kwargs (user-tunable):
  per_channel = True
pinned by method (not overridable):
  bits = 9
check_backend_contract (toy_int9):
  [PASS] quantized_linear_classes() is non-empty: ['Int9Linear']
  [PASS] Int9Linear subclasses torch.nn.Linear
  [PASS] a plain nn.Linear is not reported as quantized
  [PASS] create_quantized_linear_shell() returns a declared class, got Int9Linear
  [PASS] the shell is recognized before load_state_dict (disk offload routing)
  [PASS] create_quantized_linear() returns a declared class, got Int9Linear
  [PASS] every stored key lives under the layer name; uncovered: []
  => OK

Next, check the numerical error, the real memory saving, the dtype guard of clause (b), and the differentiability of clause (d) on a small model:

import torch
from diffsynth.core.quant import QuantizeConfig, check_differentiable


class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(256, 512)
        self.fc2 = torch.nn.Linear(512, 256, bias=False)

    def forward(self, x):
        return self.fc2(torch.nn.functional.silu(self.fc1(x)))


def footprint(model):
    return sum(t.numel() * t.element_size() for t in list(model.parameters()) + list(model.buffers()))


torch.manual_seed(0)
model = ToyModel().to(torch.bfloat16)
x = torch.randn(4, 256, dtype=torch.bfloat16)
reference = model(x)
fp_bytes = footprint(model)

QuantizeConfig(method="toy_int9_w9a16").quantize_model(model, compute_device="cpu")
print("relative error:", ((model(x) - reference).norm() / reference.norm()).item())
print(f"footprint: {fp_bytes} -> {footprint(model)} bytes ({footprint(model) / fp_bytes:.3f} of bf16)")

model.to(torch.float32)                                   # clause (b): packed dtypes must not change
print(model.fc1.weight.dtype, model.fc1.weight_msb.dtype, model.fc1.weight_scale.dtype, model.fc1.bias.dtype)

check_differentiable(model.fc1)                           # clause (d)

2 nn.Linear layers quantized (method: toy_int9_w9a16).
relative error: 0.004150390625
footprint: 525312 -> 299008 bytes (0.569 of bf16)
torch.uint8 torch.uint8 torch.float32 torch.float32
check_differentiable (Int9Linear): OK -- gradients pass through the module to its input

The measured footprint is 0.569 of bf16, slightly above 9/16 = 0.5625 because of the fp32 scales and the unquantized bias. If this ratio comes out close to 1, the packing format is not actually compressing the weights and you should revisit the storage layout in Step 1.

Inference on a Real Model: Z-Image Copy Section

Once the small-model checks pass, the backend is ready for real models — a custom backend is used exactly like a built-in method. Import the module that registers it, then pass the method to ModelConfig(quantize=...):

import torch

import my_project.toy_int9   # registers the toy_int9 backend and the toy_int9_w9a16 method
from diffsynth.core.quant import QuantizeConfig
from diffsynth.pipelines.z_image import ModelConfig, ZImagePipeline

pipe = ZImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(
            model_id="Tongyi-MAI/Z-Image-Turbo",
            origin_file_pattern="transformer/*.safetensors",
            quantize=QuantizeConfig(method="toy_int9_w9a16"),
        ),
        ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors"),
        ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
    ],
    tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
)

dit_bytes = sum(t.numel() * t.element_size() for t in list(pipe.dit.parameters()) + list(pipe.dit.buffers()))
print(f"dit weights: {dit_bytes / 1024 ** 3:.3f} GiB")

prompt = "A delicate portrait of an underwater girl, blue dress flowing, hair gently drifting, light and shadow clear, surrounded by bubbles, serene expression, exquisite details, dreamlike and beautiful."
image = pipe(prompt=prompt, seed=42, rand_device="cuda")
image.save("z_image_toy_int9.jpg")

Measured DiT weight footprint on Z-Image Turbo (the 8-step Turbo generation works normally, with no visible quality difference from bf16):

Note that peak memory and weight footprint are not the same thing: this toy materializes a temporary fp weight on every forward, so the peak saving is smaller than the storage saving. Measured on a synthetic model with 48 Linears, all weights resident on the GPU:

That temporary weight depends only on the largest single layer and does not grow with depth, so the deeper the model, the closer the peak saving gets to the weight ratio; a real backend fusing unpacking into the matmul kernel does not need it at all. To push the peak down further, stack VRAM management on top and move weights layer by layer (pass a vram_config to each ModelConfig above — measured peak drops to 2.1 GiB).

Accuracy: int9 vs int8 Copy Section

Does the extra bit actually buy accuracy? Quantize the same weight with the identical per-channel symmetric scheme at 8 and 9 bits, then compare the dequantized weight error and the layer output error. This is the general recipe for an accuracy regression on a new backend: hold everything else fixed and change only the bit width.

import torch
from my_project.toy_int9 import Int9QuantBackend, Int9WeightOnlyConfig


def quantize_int8(linear):
    """The same per-channel symmetric scheme with one bit less: codes in [-128, 127]."""
    weight = linear.weight.data
    scale = (weight.abs().amax(dim=1).float() / 127).clamp(min=1e-8)
    codes = (weight.float() / scale.unsqueeze(1)).round().clamp(-128, 127)
    return (codes * scale.unsqueeze(1)).to(weight.dtype)


def relative_error(reference, value):
    return ((value.float() - reference.float()).norm() / reference.float().norm()).item()


torch.manual_seed(0)
backend = Int9QuantBackend(Int9WeightOnlyConfig())
linear = torch.nn.Linear(2048, 2048, bias=False).to(torch.bfloat16)
fp_weight = linear.weight.data.clone()

int9_weight = backend.create_quantized_linear(linear).dequantize_weight(torch.bfloat16)
int8_weight = quantize_int8(linear)
error8, error9 = relative_error(fp_weight, int8_weight), relative_error(fp_weight, int9_weight)
print(f"weight error: int8 {error8:.6f} | int9 {error9:.6f} ({error8 / error9:.2f}x lower)")

x = torch.randn(64, 2048, dtype=torch.bfloat16)
reference = torch.nn.functional.linear(x, fp_weight)
out8 = relative_error(reference, torch.nn.functional.linear(x, int8_weight))
out9 = relative_error(reference, torch.nn.functional.linear(x, int9_weight))
print(f"output error: int8 {out8:.6f} | int9 {out9:.6f} ({out8 / out9:.2f}x lower)")

weight error: int8 0.004353 | int9 0.001937 (2.25x lower)
output error: int8 0.004947 | int9 0.002816 (1.76x lower)

This matches the theory: going from 255 to 511 levels halves the quantization step, and for uniform quantization the error is proportional to the step, so the weight error drops to roughly half (2.25x measured). The end-to-end layer output gain is smaller (1.76x) because the activations themselves are bf16 and the matmul's own rounding noise eats part of the benefit — a reminder to evaluate bit-width gains at the actual compute precision, not only on the weights.

Finally, verify clause (c): save the quantized weights, load them back into shells, and confirm both produce identical outputs.

from safetensors.torch import load_file, save_file

save_config = QuantizeConfig(method="toy_int9_w9a16")
tensors, metadata = save_config.flatten_state_dict(model.state_dict())
save_file(tensors, "toy_int9.safetensors", metadata=metadata)

loaded = ToyModel().to(torch.bfloat16)
load_config = QuantizeConfig(method="toy_int9_w9a16", load_prequantized=True)
load_config.prepare_for_prequantized_load(loaded, compute_dtype=torch.bfloat16)
loaded.load_state_dict(load_config.unflatten_state_dict(load_file("toy_int9.safetensors"), metadata), assign=True)
print("reload match:", torch.equal(loaded(x.float()), model(x.float())))

reload match: True

Combining with Disk Offload Copy Section

Disk offload, part of VRAM management, places the strictest demands on a quantization backend: the resident model keeps only meta shells, and each layer's tensors are streamed back from disk at forward time and dropped right after. It relies on two things:

• It only supports pre-quantized checkpoints, so load_prequantized=True is required and prepare_for_prequantized_load must first swap the target layers for shells.
• Which tensors a layer needs is resolved by a prefix scan over the checkpoint keys using the layer's dotted name, and the result is loaded with a strict load_state_dict(assign=True). So the only requirement on a backend is that every tensor lives under {layer_name}. — flat siblings like layer.weight_scale and nested quant state like bnb's both work. A missing or extra key raises instead of silently loading a corrupted layer.

/* Detailed source-code truncated for AI context efficiency. */

Measured on the same ToyModel (torch_dtype=None on DiskMap is essential — it guarantees the packed tensors are not re-typed while being read):

2 nn.Linear layers replaced for loading the pre-quantized checkpoint (method: toy_int9_w9a16).
fc1: ['fc1.bias', 'fc1.weight', 'fc1.weight_msb', 'fc1.weight_scale']
fc2: ['fc2.weight', 'fc2.weight_msb', 'fc2.weight_scale']
resident 299008 bytes -> in memory after disk offload 0 bytes
output matches: True | repeatable: True

Each layer's weight / weight_msb / weight_scale / bias is correctly attributed to that layer, the resident footprint drops to 0 bytes (everything is a meta shell), the output is bit-identical to the resident quantized model, and repeated forwards stay stable — so rebuilding shells and streaming from disk has no side effects.

On a real model, use the standard workflows from Model Quantization for end-to-end validation: pass QuantizeConfig(method="toy_int9_w9a16") to ModelConfig(quantize=...) for online-quantized inference, save the quantized weights with save_quantized_model and load them back after registering the hash, and inject LoRA into the quantized model for training.

Integration Checklist Copy Chapter

• The packing format really shrinks the weights: the measured footprint ratio should be close to the theoretical bit-width ratio, not close to 1.
• The accuracy gain is verified: compared against the same scheme with one bit less, the error really goes down; otherwise precision is being lost somewhere in the dequantization path.
• The quantized Linear subclasses torch.nn.Linear, and all state_dict keys live under the layer name.
• _apply guards the dtype of every packed tensor and quant state.
• capabilities() matches reality: declaring is_serializable requires a round-tripping state dict, and declaring is_differentiable requires passing check_differentiable.
• create_quantized_linear honors compute_device / model_device, so layer-by-layer streaming quantization works.
• Disk offload works: the shell is built on meta and cheap to rebuild, every stored tensor lives under the layer's dotted name, and unflatten_state_dict tolerates being called with a single layer's subdict plus whole-file metadata.
• When depending on a third-party library, validate_environment() gives a clear installation hint and project_url points at the upstream project.
• check_backend_contract passes completely.

En Developer Guide Integrating Your Model (`docs/en/Developer_Guide/Integrating_Your_Model.md`) Copy Section

Integrating Model Architecture

This document introduces how to integrate models into the DiffSynth-Studio framework for use by modules such as Pipeline.

Step 1: Integrate Model Architecture Code Copy Chapter

All model architecture implementations in DiffSynth-Studio are unified in diffsynth/models. Each .py code file implements a model architecture, and all models are loaded through ModelPool in diffsynth/models/model_loader.py. When integrating new model architectures, please create a new .py file under this path.

diffsynth/models/
├── general_modules.py
├── model_loader.py
├── qwen_image_controlnet.py
├── qwen_image_dit.py
├── qwen_image_text_encoder.py
├── qwen_image_vae.py
└── ...

In most cases, we recommend integrating models in native PyTorch code form, with the model architecture class directly inheriting from torch.nn.Module, for example:

import torch

class NewDiffSynthModel(torch.nn.Module):
    def __init__(self, dim=1024):
        super().__init__()
        self.linear = torch.nn.Linear(dim, dim)
        self.activation = torch.nn.Sigmoid()
    
    def forward(self, x):
        x = self.linear(x)
        x = self.activation(x)
        return x

If the model architecture implementation contains additional dependencies, we strongly recommend removing them, otherwise this will cause heavy package dependency issues. In our existing models, Qwen-Image's Blockwise ControlNet is integrated in this way. The code is lightweight, please refer to diffsynth/models/qwen_image_controlnet.py.

If the model has been integrated by Huggingface Library (transformers, diffusers, etc.), we can integrate the model in a simpler way:

Integrating Huggingface Library Style Model Architecture Code

The loading method for these models in Huggingface Library is:

from transformers import XXX_Model

model = XXX_Model.from_pretrained("path_to_your_model")

DiffSynth-Studio does not support loading models through from_pretrained because this conflicts with VRAM management and other functions. Please rewrite the model architecture in the following format:

import torch

class DiffSynth_XXX_Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        from transformers import XXX_Config, XXX_Model
        config = XXX_Config(**{
            "architectures": ["XXX_Model"],
            "other_configs": "Please copy and paste the other configs here.",
        })
        self.model = XXX_Model(config)
        
    def forward(self, x):
        outputs = self.model(x)
        return outputs

Where XXX_Config is the Config class corresponding to the model. For example, the Config class for Qwen2_5_VLModel is Qwen2_5_VLConfig, which can be found by consulting its source code. The content inside Config can usually be found in the config.json file in the model library. DiffSynth-Studio will not read the config.json file, so the content needs to be copied and pasted into the code.

In rare cases, version updates of transformers and diffusers may cause some models to be unable to import. Therefore, if possible, we still recommend using the model integration method in Step 1.1.

In our existing models, Qwen-Image's Text Encoder is integrated in this way. The code is lightweight, please refer to diffsynth/models/qwen_image_text_encoder.py.

Step 2: Model File Format Conversion Copy Chapter

Due to the variety of model file formats provided by developers in the open-source community, we sometimes need to convert model file formats to form correctly formatted state dict. This is common in the following situations:

• Model files built by different code libraries, for example Wan-AI/Wan2.1-T2V-1.3B and Wan-AI/Wan2.1-T2V-1.3B-Diffusers.
• Models modified during integration, for example, the Text Encoder of Qwen/Qwen-Image adds a model. prefix in diffsynth/models/qwen_image_text_encoder.py.
• Model files containing multiple models, for example, the VACE Adapter and base DiT model of Wan-AI/Wan2.1-VACE-14B are mixed and stored in the same set of model files.

In our development philosophy, we hope to respect the wishes of model authors as much as possible. If we repackage the model files, for example Comfy-Org/Qwen-Image_ComfyUI, although we can call the model more conveniently, traffic (model page views and downloads, etc.) will be directed elsewhere, and the original author of the model will also lose the power to delete the model. Therefore, we have added the diffsynth/utils/state_dict_converters module to the framework for file format conversion during model loading.

This part of logic is very simple. Taking Qwen-Image's Text Encoder as an example, only 10 lines of code are needed:

def QwenImageTextEncoderStateDictConverter(state_dict):
    state_dict_ = {}
    for k in state_dict:
        v = state_dict[k]
        if k.startswith("visual."):
            k = "model." + k
        elif k.startswith("model."):
            k = k.replace("model.", "model.language_model.")
        state_dict_[k] = v
    return state_dict_

Step 3: Writing Model Config Copy Chapter

Model Config is located in diffsynth/configs/model_configs.py, used to identify model types and load them. The following fields need to be filled in:

• model_hash: Model file hash value, which can be obtained through the hash_model_file function. This hash value is only related to the keys and tensor shapes in the model file's state dict, and is unrelated to other information in the file.
• model_name: Model name, used for Pipeline to identify the required model. If different structured models play the same role in Pipeline, the same model_name can be used. When integrating new models, just ensure that model_name is different from other existing functional models. The corresponding model is fetched through model_name in the Pipeline's from_pretrained.
• model_class: Model architecture import path, pointing to the model architecture class implemented in Step 1, for example diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder.
• state_dict_converter: Optional parameter. If model file format conversion is needed, the import path of the model conversion logic needs to be filled in, for example diffsynth.utils.state_dict_converters.qwen_image_text_encoder.QwenImageTextEncoderStateDictConverter.
• extra_kwargs: Optional parameter. If additional parameters need to be passed when initializing the model, these parameters need to be filled in. For example, models DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny and DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Inpaint both adopt the QwenImageBlockWiseControlNet structure in diffsynth/models/qwen_image_controlnet.py, but the latter also needs additional configuration additional_in_dim=4. Therefore, this configuration information needs to be filled in the extra_kwargs field.

We provide a piece of code to quickly understand how models are loaded through this configuration information:

from diffsynth.core import hash_model_file, load_state_dict, skip_model_initialization
from diffsynth.models.qwen_image_text_encoder import QwenImageTextEncoder
from diffsynth.utils.state_dict_converters.qwen_image_text_encoder import QwenImageTextEncoderStateDictConverter
import torch

model_hash = "8004730443f55db63092006dd9f7110e"
model_name = "qwen_image_text_encoder"
model_class = QwenImageTextEncoder
state_dict_converter = QwenImageTextEncoderStateDictConverter
extra_kwargs = {}

model_path = [
    "models/Qwen/Qwen-Image/text_encoder/model-00001-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00002-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00003-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00004-of-00004.safetensors",
]
if hash_model_file(model_path) == model_hash:
    with skip_model_initialization():
        model = model_class(**extra_kwargs)
    state_dict = load_state_dict(model_path, torch_dtype=torch.bfloat16, device="cuda")
    state_dict = state_dict_converter(state_dict)
    model.load_state_dict(state_dict, assign=True)
    print("Done!")

Q: The logic of the above code looks very simple, why is this part of code in DiffSynth-Studio extremely complex?

A: Because we provide aggressive VRAM management functions that are coupled with the model loading logic, this leads to the complexity of the framework structure. We have tried our best to simplify the interface exposed to developers.

The model_hash in diffsynth/configs/model_configs.py is not uniquely existing. Multiple models may exist in the same model file. For this situation, please use multiple model Configs to load each model separately, and write the corresponding state_dict_converter to separate the parameters required by each model.

Step 4: Verifying Whether the Model Can Be Recognized and Loaded Copy Chapter

After model integration, the following code can be used to verify whether the model can be correctly recognized and loaded. The following code will attempt to load the model into memory:

from diffsynth.models.model_loader import ModelPool

model_pool = ModelPool()
model_pool.auto_load_model(
    [
        "models/Qwen/Qwen-Image/text_encoder/model-00001-of-00004.safetensors",
        "models/Qwen/Qwen-Image/text_encoder/model-00002-of-00004.safetensors",
        "models/Qwen/Qwen-Image/text_encoder/model-00003-of-00004.safetensors",
        "models/Qwen/Qwen-Image/text_encoder/model-00004-of-00004.safetensors",
    ],
)

If the model can be recognized and loaded, you will see the following output:

Loading models from: [
    "models/Qwen/Qwen-Image/text_encoder/model-00001-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00002-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00003-of-00004.safetensors",
    "models/Qwen/Qwen-Image/text_encoder/model-00004-of-00004.safetensors"
]
Loaded model: {
    "model_name": "qwen_image_text_encoder",
    "model_class": "diffsynth.models.qwen_image_text_encoder.QwenImageTextEncoder",
    "extra_kwargs": null
}

Step 5: Writing Model VRAM Management Scheme Copy Chapter

DiffSynth-Studio supports complex VRAM management. See Enabling VRAM Management for details.

En Developer Guide Training Diffusion Models (`docs/en/Developer_Guide/Training_Diffusion_Models.md`) Copy Section

Integrating Model Training

After integrating models and implementing Pipeline, the next step is to integrate model training functionality.

Training-Inference Consistent Pipeline Modification Copy Chapter

To ensure strict consistency between training and inference processes, we will use most of the inference code during training, but still need to make minor modifications.

First, add extra logic during inference to switch the image-to-image/video-to-video logic based on the scheduler state. Taking Qwen-Image as an example:

class QwenImageUnit_InputImageEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("input_image", "noise", "tiled", "tile_size", "tile_stride"),
            output_params=("latents", "input_latents"),
            onload_model_names=("vae",)
        )

    def process(self, pipe: QwenImagePipeline, input_image, noise, tiled, tile_size, tile_stride):
        if input_image is None:
            return {"latents": noise, "input_latents": None}
        pipe.load_models_to_device(['vae'])
        image = pipe.preprocess_image(input_image).to(device=pipe.device, dtype=pipe.torch_dtype)
        input_latents = pipe.vae.encode(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
        if pipe.scheduler.training:
            return {"latents": noise, "input_latents": input_latents}
        else:
            latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0])
            return {"latents": latents, "input_latents": input_latents}

Then, enable Gradient Checkpointing in model_fn, which will significantly reduce the VRAM required for training at the cost of computational speed. This is not mandatory, but we strongly recommend doing so.

Taking Qwen-Image as an example, before modification:

text, image = block(
    image=image,
    text=text,
    temb=conditioning,
    image_rotary_emb=image_rotary_emb,
    attention_mask=attention_mask,
)

After modification:

from ..core import gradient_checkpoint_forward

text, image = gradient_checkpoint_forward(
    block,
    use_gradient_checkpointing,
    use_gradient_checkpointing_offload,
    image=image,
    text=text,
    temb=conditioning,
    image_rotary_emb=image_rotary_emb,
    attention_mask=attention_mask,
)

Writing Training Scripts Copy Chapter

DiffSynth-Studio does not strictly encapsulate the training framework, but exposes the script content to developers. This approach makes it more convenient to modify training scripts to implement additional functions. Developers can refer to existing training scripts, such as examples/qwen_image/model_training/train.py, for modification to adapt to new model training.

En Diffusion Templates Introducing Diffusion Templates (`docs/en/Diffusion_Templates/Introducing_Diffusion_Templates.md`) Copy Section

Diffusion Templates

Diffusion Templates is a controllable generation plugin framework for Diffusion models in DiffSynth-Studio, providing additional controllable generation capabilities for base models.

• Open-source code: DiffSynth-Studio
• Technical report: arXiv
• Project page: GitHub
• Documentation reference
  • Introduction to Diffusion Templates: English Version、中文版
  • Detailed Architecture of Diffusion Templates: English Version、中文版
  • Template Model Inference: English Version、中文版
  • Template Model Training: English Version、中文版
• Online demo: ModelScope
• Model collection: ModelScope、ModelScope International、HuggingFace

• Dataset: ModelScope、ModelScope International、HuggingFace

Model Performance Overview Copy Chapter

• Super-Resolution + Sharpness Enhancement: Generate ultra-high-resolution images

• Structure Control + Aesthetic Alignment + Sharpness Enhancement: Fully-equipped ControlNet

• Structure Control + Image Editing + Color Adjustment: Artistic Style Creation at Will

• Brightness Control + Image Editing + Local Redrawing: Cross-dimensional Elements in Images

En Diffusion Templates Template Model Inference (`docs/en/Diffusion_Templates/Template_Model_Inference.md`) Copy Section

Template Model Inference

Enabling Template Models on Base Model Pipelines Copy Chapter

Using the base model black-forest-labs/FLUX.2-klein-base-4B as an example, when generating images using only the base model:

from diffsynth.diffusion.template import TemplatePipeline
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
import torch

# Load base model
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/"),
)
# Generate an image
image = pipe(
    prompt="a cat",
    seed=0, cfg_scale=4,
    height=1024, width=1024,
)
image.save("image.png")

The Template model DiffSynth-Studio/Template-KleinBase4B-Brightness can control image brightness during generation. Through the TemplatePipeline model, it can be loaded from ModelScope (via ModelConfig(model_id="xxx/xxx")) or from a local path (via ModelConfig(path="xxx")). Inputting scale=0.8 increases image brightness. Note that in the code, input parameters for pipe must be transferred to template_pipeline, and template_inputs should be added.

# Load Template model
template_pipeline = TemplatePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Brightness")
    ],
)
# Generate an image
image = template_pipeline(
    pipe,
    prompt="a cat",
    seed=0, cfg_scale=4,
    height=1024, width=1024,
    template_inputs=[{"scale": 0.8}],
)
image.save("image_0.8.png")

CFG Enhancement for Template Models Copy Chapter

Template models can enable CFG (Classifier-Free Guidance) to make control effects more pronounced. For example, with the model DiffSynth-Studio/Template-KleinBase4B-Brightness, adding negative_template_inputs to the TemplatePipeline input parameters and setting its scale to 0.5 will generate images with more noticeable brightness variations by contrasting both sides.

# Generate an image with CFG
image = template_pipeline(
    pipe,
    prompt="a cat",
    seed=0, cfg_scale=4,
    height=1024, width=1024,
    template_inputs=[{"scale": 0.8}],
    negative_template_inputs=[{"scale": 0.5}],
)
image.save("image_0.8_cfg.png")

Low VRAM Support Copy Chapter

Template models currently do not support the main framework's VRAM management, but lazy loading can be used - loading Template models only when needed for inference. This significantly reduces VRAM requirements when enabling multiple Template models, with peak VRAM usage being that of a single Template model. Add parameter lazy_loading=True to enable.

template_pipeline = TemplatePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Brightness")
    ],
    lazy_loading=True,
)

The base model's Pipeline and Template Pipeline are completely independent and can enable VRAM management on demand.

When Template model outputs contain LoRA in Template Cache, you need to enable VRAM management for the base model's Pipeline or enable LoRA hot loading (using the code below), otherwise LoRA weights will be fused repeatedly.

pipe.dit = pipe.enable_lora_hot_loading(pipe.dit)

Enabling Multiple Template Models Copy Chapter

TemplatePipeline can load multiple Template models. During inference, use model_id in template_inputs to distinguish inputs for each Template model.

After enabling VRAM management for the base model's Pipeline and lazy loading for Template Pipeline, you can load any number of Template models.

/* Detailed source-code truncated for AI context efficiency. */

Super-Resolution + Sharpness Enhancement Copy Section

Combining DiffSynth-Studio/Template-KleinBase4B-Upscaler and DiffSynth-Studio/Template-KleinBase4B-Sharpness can upscale blurry images while improving detail clarity.

image = template(
    pipe,
    prompt="A cat is sitting on a stone.",
    seed=0, cfg_scale=4, num_inference_steps=50,
    template_inputs = [
        {
            "model_id": 3,
            "image": Image.open("data/examples/templates/image_lowres_100.jpg"),
            "prompt": "A cat is sitting on a stone.",
        },
        {
            "model_id": 5,
            "scale": 1,
        },
    ],
    negative_template_inputs = [
        {
            "model_id": 3,
            "image": Image.open("data/examples/templates/image_lowres_100.jpg"),
            "prompt": "",
        },
        {
            "model_id": 5,
            "scale": 0,
        },
    ],
)
image.save("image_Upscaler_Sharpness.png")

Structure Control + Aesthetic Alignment + Sharpness Enhancement Copy Section

DiffSynth-Studio/Template-KleinBase4B-ControlNet controls composition, DiffSynth-Studio/Template-KleinBase4B-Aesthetic fills in details, and DiffSynth-Studio/Template-KleinBase4B-Sharpness ensures clarity. Combining these three Template models produces exquisite images.

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 = [
        {
            "model_id": 1,
            "image": Image.open("data/examples/templates/image_depth.jpg"),
            "prompt": "A cat is sitting on a stone, bathed in bright sunshine.",
        },
        {
            "model_id": 7,
            "lora_ids": list(range(1, 180, 2)),
            "lora_scales": 2.0,
            "merge_type": "mean",
        },
        {
            "model_id": 5,
            "scale": 0.8,
        },
    ],
    negative_template_inputs = [
        {
            "model_id": 1,
            "image": Image.open("data/examples/templates/image_depth.jpg"),
            "prompt": "",
        },
        {
            "model_id": 7,
            "lora_ids": list(range(1, 180, 2)),
            "lora_scales": 2.0,
            "merge_type": "mean",
        },
        {
            "model_id": 5,
            "scale": 0,
        },
    ],
)
image.save("image_Controlnet_Aesthetic_Sharpness.png")

Structure Control + Image Editing + Color Adjustment Copy Section

DiffSynth-Studio/Template-KleinBase4B-ControlNet controls composition, DiffSynth-Studio/Template-KleinBase4B-Edit preserves original image details like fur texture, and DiffSynth-Studio/Template-KleinBase4B-SoftRGB controls color tones, creating an artistic masterpiece.

image = template(
    pipe,
    prompt="A cat is sitting on a stone. Colored ink painting.",
    seed=0, cfg_scale=4, num_inference_steps=50,
    template_inputs = [
        {
            "model_id": 1,
            "image": Image.open("data/examples/templates/image_depth.jpg"),
            "prompt": "A cat is sitting on a stone. Colored ink painting.",
        },
        {
            "model_id": 2,
            "image": Image.open("data/examples/templates/image_reference.jpg"),
            "prompt": "Convert the image style to colored ink painting.",
        },
        {
            "model_id": 4,
            "R": 0.9,
            "G": 0.5,
            "B": 0.3,
        },
    ],
    negative_template_inputs = [
        {
            "model_id": 1,
            "image": Image.open("data/examples/templates/image_depth.jpg"),
            "prompt": "",
        },
        {
            "model_id": 2,
            "image": Image.open("data/examples/templates/image_reference.jpg"),
            "prompt": "",
        },
    ],
)
image.save("image_Controlnet_Edit_SoftRGB.png")

Brightness Control + Image Editing + Local Redrawing Copy Section

DiffSynth-Studio/Template-KleinBase4B-Brightness generates bright scenes, DiffSynth-Studio/Template-KleinBase4B-Edit references original image layout, and DiffSynth-Studio/Template-KleinBase4B-Inpaint keeps background unchanged, generating cross-dimensional content.

image = template(
    pipe,
    prompt="A cat is sitting on a stone. Flat anime style.",
    seed=0, cfg_scale=4, num_inference_steps=50,
    template_inputs = [
        {
            "model_id": 0,
            "scale": 0.6,
        },
        {
            "model_id": 2,
            "image": Image.open("data/examples/templates/image_reference.jpg"),
            "prompt": "Convert the image style to flat anime style.",
        },
        {
            "model_id": 6,
            "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 = [
        {
            "model_id": 0,
            "scale": 0.5,
        },
        {
            "model_id": 2,
            "image": Image.open("data/examples/templates/image_reference.jpg"),
            "prompt": "",
        },
        {
            "model_id": 6,
            "image": Image.open("data/examples/templates/image_reference.jpg"),
            "mask": Image.open("data/examples/templates/image_mask_1.jpg"),
        },
    ],
)
image.save("image_Brightness_Edit_Inpaint.png")

En Diffusion Templates Template Model Training (`docs/en/Diffusion_Templates/Template_Model_Training.md`) Copy Section

Template Model Training

DiffSynth-Studio currently provides comprehensive Template training support for black-forest-labs/FLUX.2-klein-base-4B, with more model adaptations coming soon.

Continuing Training from Pretrained Models Copy Chapter

To continue training from our pretrained models, refer to the table in FLUX.2 to find the corresponding training script.

Building New Template Models Copy Chapter

Template Model Component Format Copy Section

A Template model binds to a model repository (or local folder) containing a code file model.py as the entry point. Here's the template for model.py:

import torch

class CustomizedTemplateModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    @torch.no_grad()
    def process_inputs(self, xxx, **kwargs):
        yyy = xxx
        return {"yyy": yyy}

    def forward(self, yyy, **kwargs):
        zzz = yyy
        return {"zzz": zzz}

class DataProcessor:
    def __call__(self, www, **kwargs):
        xxx = www
        return {"xxx": xxx}

TEMPLATE_MODEL = CustomizedTemplateModel
TEMPLATE_MODEL_PATH = "model.safetensors"
TEMPLATE_DATA_PROCESSOR = DataProcessor

During Template model inference, Template Input passes through TEMPLATE_MODEL's process_inputs and forward to generate Template Cache.

flowchart LR;
    i@{shape: text, label: "Template Input"}-->p[process_inputs];
    subgraph TEMPLATE_MODEL
        p[process_inputs]-->f[forward]
    end
    f[forward]-->c@{shape: text, label: "Template Cache"};

During Template model training, Template Input comes from the dataset through TEMPLATE_DATA_PROCESSOR.

flowchart LR;
    d@{shape: text, label: "Dataset"}-->dp[TEMPLATE_DATA_PROCESSOR]-->p[process_inputs];
    subgraph TEMPLATE_MODEL
        p[process_inputs]-->f[forward]
    end
    f[forward]-->c@{shape: text, label: "Template Cache"};

`TEMPLATE_MODEL`

TEMPLATE_MODEL implements the Template model logic, inheriting from torch.nn.Module with required process_inputs and forward methods. These two methods form the complete Template model inference process, split into two stages to better support two-stage split training.

• process_inputs must use @torch.no_grad() for gradient-free computation
• forward must contain all gradient computations required for training

Both methods should accept **kwargs for compatibility. Reserved parameters include:

• To interact with the base model Pipeline (e.g., call text encoder), add pipe parameter to method inputs
• To enable Gradient Checkpointing, add use_gradient_checkpointing and use_gradient_checkpointing_offload to forward inputs
• Multiple Template models use model_id to distinguish Template Inputs - do not use this field in method parameters

`TEMPLATE_MODEL_PATH` (Optional)

TEMPLATE_MODEL_PATH specifies the relative path to pretrained weights. For example:

TEMPLATE_MODEL_PATH = "model.safetensors"

For multi-file models:

TEMPLATE_MODEL_PATH = [
    "model-00001-of-00003.safetensors",
    "model-00002-of-00003.safetensors",
    "model-00003-of-00003.safetensors",
]

Set to None for random initialization:

TEMPLATE_MODEL_PATH = None

`TEMPLATE_DATA_PROCESSOR` (Optional)

To train Template models with DiffSynth-Studio, datasets should contain template_inputs fields in metadata.json. These fields pass through TEMPLATE_DATA_PROCESSOR to generate inputs for Template model methods.

For example, the brightness control model DiffSynth-Studio/Template-KleinBase4B-Brightness takes scale as input:

[
    {
        "image": "images/image_1.jpg",
        "prompt": "a cat",
        "template_inputs": {"scale": 0.2}
    },
    {
        "image": "images/image_2.jpg",
        "prompt": "a dog",
        "template_inputs": {"scale": 0.6}
    }
]

class DataProcessor:
    def __call__(self, scale, **kwargs):
        return {"scale": scale}

TEMPLATE_DATA_PROCESSOR = DataProcessor

Or calculate scale from image paths:

[
    {
        "image": "images/image_1.jpg",
        "prompt": "a cat",
        "template_inputs": {"image": "/path/to/your/dataset/images/image_1.jpg"}
    }
]

class DataProcessor:
    def __call__(self, image, **kwargs):
        image = Image.open(image)
        image = np.array(image)
        return {"scale": image.astype(np.float32).mean() / 255}

TEMPLATE_DATA_PROCESSOR = DataProcessor

Training Template Models Copy Section

A Template model is "trainable" if its Template Cache variables are fully decoupled from the base model Pipeline - these variables should reach model_fn without participating in any Pipeline Unit calculations.

For training with black-forest-labs/FLUX.2-klein-base-4B, use these training script parameters:

• --extra_inputs: Additional inputs. Use template_inputs for text-to-image models, edit_image,template_inputs for image editing models
• --template_model_id_or_path: Template model ID or local path (use : suffix for ModelScope IDs, e.g., "DiffSynth-Studio/Template-KleinBase4B-Brightness:")
• --remove_prefix_in_ckpt: State dict prefix to remove when saving models (use "pipe.template_model.")
• --trainable_models: Trainable components (use "template_model" for full model, or "template_model.xxx,template_model.yyy" for specific components)

Example training script:

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 "examples/flux2/model_training/scripts/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_example" \
  --trainable_models "template_model" \
  --use_gradient_checkpointing \
  --find_unused_parameters

Interacting with Base Model Pipeline Components Copy Section

Template models can interact with base model Pipelines. For example, using the text encoder:

class CustomizedTemplateModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.xxx = xxx()

    @torch.no_grad()
    def process_inputs(self, text, pipe, **kwargs):
        input_ids = pipe.tokenizer(text)
        text_emb = pipe.text_encoder(input_ids)
        return {"text_emb": text_emb}

    def forward(self, text_emb, pipe, **kwargs):
        kv_cache = self.xxx(text_emb)
        return {"kv_cache": kv_cache}

TEMPLATE_MODEL = CustomizedTemplateModel

Using Non-Trainable Components Copy Section

For models with pretrained components:

class CustomizedTemplateModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.image_encoder = XXXEncoder.from_pretrained(xxx)
        self.mlp = MLP()

    @torch.no_grad()
    def process_inputs(self, image, **kwargs):
        emb = self.image_encoder(image)
        return {"emb": emb}

    def forward(self, emb, **kwargs):
        kv_cache = self.mlp(emb)
        return {"kv_cache": kv_cache}

TEMPLATE_MODEL = CustomizedTemplateModel

Set --trainable_models template_model.mlp to train only the MLP component.

Training on Low VRAM Devices Copy Section

The framework supports splitting Template model training into two stages: the first stage performs gradient-free computation, and the second stage performs gradient updates. For more information, refer to the documentation: Two-stage Split Training. Here's a sample script:

/* Detailed source-code truncated for AI context efficiency. */

Two-stage split training can reduce VRAM requirements and improve training speed. The training process is lossless in precision, but requires significant disk space for storing cache files.

To further reduce VRAM requirements, you can enable fp8 precision by adding the parameters --fp8_models "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" and --fp8_models "black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors" to the two-stage training. Note that fp8 precision can only be enabled on non-trainable model components and introduces minor errors.

Uploading Template Models Copy Section

After training, follow these steps to upload Template models to ModelScope for wider distribution.

1. Set model path in model.py:

TEMPLATE_MODEL_PATH = "model.safetensors"

2. Upload using ModelScope CLI:

modelscope upload user_name/your_model_id /path/to/your/model.py model.py --token ms-xxx

3. Package model files:

from diffsynth.diffusion.template import load_template_model, load_state_dict
from safetensors.torch import save_file
import torch

model = load_template_model("path/to/your/template/model", torch_dtype=torch.bfloat16, device="cpu")
state_dict = load_state_dict("path/to/your/ckpt/epoch-1.safetensors", torch_dtype=torch.bfloat16, device="cpu")
state_dict.update(model.state_dict())
save_file(state_dict, "model.safetensors")

4. Upload model file:

modelscope upload user_name/your_model_id /path/to/your/model/epoch-1.safetensors model.safetensors --token ms-xxx

5. Verify inference:

from diffsynth.diffusion.template import TemplatePipeline
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
import torch

# Load base model
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/"),
)

# Load Template model
template_pipeline = TemplatePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(model_id="user_name/your_model_id")
    ],
)

# Generate image
image = template_pipeline(
    pipe,
    prompt="a cat",
    seed=0, cfg_scale=4,
    height=1024, width=1024,
    template_inputs=[{xxx}],
)
image.save("image.png")

---

### En Diffusion Templates Understanding Diffusion Templates (`docs/en/Diffusion_Templates/Understanding_Diffusion_Templates.md`)

# Diffusion Templates Architecture Details

The Diffusion Templates framework is a controllable generation plugin framework in DiffSynth-Studio that provides additional controllable generation capabilities for Diffusion models.

## Framework Structure

The Diffusion Templates framework structure is shown below:

```mermaid
flowchart TD;
    subgraph Template Pipeline
        si@{shape: text, label: "Template Input"}-->i1@{shape: text, label: "Template Input 1"};
        si@{shape: text, label: "Template Input"}-->i2@{shape: text, label: "Template Input 2"};
        si@{shape: text, label: "Template Input"}-->i3@{shape: text, label: "Template Input 3"};
        i1@{shape: text, label: "Template Input 1"}-->m1[Template Model 1]-->c1@{shape: text, label: "Template Cache 1"};
        i2@{shape: text, label: "Template Input 2"}-->m2[Template Model 2]-->c2@{shape: text, label: "Template Cache 2"};
        i3@{shape: text, label: "Template Input 3"}-->m3[Template Model 3]-->c3@{shape: text, label: "Template Cache 3"};
        c1-->c@{shape: text, label: "Template Cache"};
        c2-->c;
        c3-->c;
    end
    i@{shape: text, label: "Model Input"}-->m[Diffusion Pipeline]-->o@{shape: text, label: "Model Output"};
    c-->m;

The framework contains these module designs:

• Template Input: Template model input. Format: Python dictionary with fields determined by each Template model (e.g., {"scale": 0.8})
• Template Model: Template model, loadable from ModelScope (ModelConfig(model_id="xxx/xxx")) or local path (ModelConfig(path="xxx"))
• Template Cache: Template model output. Format: Python dictionary with fields matching base model Pipeline input parameters
• Template Pipeline: Module for managing multiple Template models. Handles model loading and cache integration

When the Diffusion Templates framework is disabled, base model components (Text Encoder, DiT, VAE) are loaded into the Diffusion Pipeline. Model Input (prompt, height, width) produces Model Output (e.g., images).

When enabled, Template models are loaded into the Template Pipeline. The Template Pipeline outputs Template Cache (a subset of Diffusion Pipeline input parameters) for subsequent processing in the Diffusion Pipeline. This enables controllable generation by intercepting part of the Diffusion Pipeline's input parameters.

Model Capability Medium Copy Chapter

Template Cache is defined as a subset of Diffusion Pipeline input parameters, ensuring framework generality. We restrict Template model inputs to only be Diffusion Pipeline parameters. The KV-Cache is particularly suitable as a Diffusion medium:

• Proven effective in LLM Skills (prompts are converted to KV-Cache)
• Has "high permission" in Diffusion models - can directly control image generation
• Supports sequence-level concatenation for multiple Template models
• Requires minimal development (add pipeline parameter and integrate to model)

Other potential Template mediums:

• Residual: Used in ControlNet for point-to-point control, but has resolution limitations and potential conflicts when merging
• LoRA: Treated as input parameters rather than model components

Currently, we only support KV-Cache and LoRA as Template Cache mediums in FLUX.2 Pipeline, with plans to support more models and mediums in the future.

Template Model Format Copy Chapter

A Template model has this structure:

Template_Model
├── model.py
└── model.safetensors

Where model.py is the entry point and model.safetensors contains model weights. For implementation details, see Template Model Training or existing Template models.

En Model Details ACE Step (`docs/en/Model_Details/ACE-Step.md`) Copy Section

ACE-Step

ACE-Step 1.5 is an open-source music generation model based on DiT architecture, supporting text-to-music, audio cover, repainting and other functionalities, running efficiently on consumer-grade hardware.

Installation Copy Chapter

Before performing model inference and training, please install DiffSynth-Studio first.

git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
pip install -e .

For more information on installation, please refer to Setup Dependencies.

Quick Start Copy Chapter

Running the following code will load the ACE-Step/Ace-Step1.5 model for inference. VRAM management is enabled, the framework automatically controls parameter loading based on available VRAM, requiring a minimum of 3GB VRAM.

/* Detailed source-code truncated for AI context efficiency. */

Model Overview Copy Chapter

Model Inference Copy Chapter

The model is loaded via AceStepPipeline.from_pretrained, see Loading Models for details.

The input parameters for AceStepPipeline inference include:

• prompt: Text description of the music.
• cfg_scale: Classifier-free guidance scale, defaults to 1.0.
• lyrics: Lyrics text.
• task_type: Task type,可选 values include "text2music" (text-to-music), "cover" (audio cover), "repaint" (repainting), defaults to "text2music".
• reference_audios: List of reference audio tensors for timbre reference.
• src_audio: Source audio tensor for cover or repaint tasks.
• denoising_strength: Denoising strength, controlling how much the output is influenced by source audio, defaults to 1.0.
• audio_cover_strength: Audio cover step ratio, controlling how many steps use cover condition in cover tasks, defaults to 1.0.
• audio_code_string: Input audio code string for cover tasks with discrete audio codes.
• repainting_ranges: List of repainting time ranges (tuples of floats, in seconds) for repaint tasks.
• repainting_strength: Repainting intensity, controlling the degree of change in repainted areas, defaults to 1.0.
• duration: Audio duration in seconds, defaults to 60.
• bpm: Beats per minute, defaults to 100.
• keyscale: Musical key scale, defaults to "B minor".
• timesignature: Time signature, defaults to "4".
• vocal_language: Vocal language, defaults to "unknown".
• seed: Random seed.
• rand_device: Device for noise generation, defaults to "cpu".
• num_inference_steps: Number of inference steps, defaults to 8.
• shift: Timestep shift parameter for the scheduler, defaults to 3.0.

Model Training Copy Chapter

Models in the ace_step series are trained uniformly via examples/ace_step/model_training/train.py. The script parameters include:

• General Training Parameters
  • Dataset Configuration
    • --dataset_base_path: Root directory of the dataset.
    • --dataset_metadata_path: Path to the dataset metadata file.
    • --dataset_repeat: Number of dataset repeats per epoch.
    • --dataset_num_workers: Number of processes per DataLoader.
    • --data_file_keys: Field names to load from metadata, typically paths to image or video files, separated by ,.
  • Model Loading Configuration
    • --model_paths: Paths to load models from, in JSON format.
    • --model_id_with_origin_paths: Model IDs with original paths, separated by commas.
    • --extra_inputs: Additional input parameters required by the model Pipeline, separated by ,.
    • --fp8_models: Models to load in FP8 format, currently only supported for models whose parameters are not updated by gradients.
    • --quant_options: Dynamically quantize loaded models. Semicolon-separated entries, each <model_string>:<method>[/<exclude_modules>], where <model_string> matches an entry in --model_paths/--model_id_with_origin_paths, method is a registered method (e.g. bitsandbytes_nf4), and exclude_modules optionally lists layers kept in full precision.
  • Basic Training Configuration
    • --learning_rate: Learning rate.
    • --num_epochs: Number of epochs.
    • --trainable_models: Trainable models, e.g., dit, vae, text_encoder.
    • --find_unused_parameters: Whether unused parameters exist in DDP training.
    • --weight_decay: Weight decay magnitude.
    • --task: Training task, defaults to sft.
  • Output Configuration
    • --output_path: Path to save the model.
    • --remove_prefix_in_ckpt: Remove prefix in the model's state dict.
    • --save_steps: Interval in training steps to save the model.
  • LoRA Configuration
    • --lora_base_model: Which model to add LoRA to.
    • --lora_target_modules: Which layers to add LoRA to.
    • --lora_rank: Rank of LoRA.
    • --lora_checkpoint: Path to LoRA checkpoint.
    • --preset_lora_path: Path to preset LoRA checkpoint for LoRA differential training.
    • --preset_lora_model: Which model to integrate preset LoRA into, e.g., dit.
  • Gradient Configuration
    • --use_gradient_checkpointing: Whether to enable gradient checkpointing.
    • --use_gradient_checkpointing_offload: Whether to offload gradient checkpointing to CPU memory.
    • --gradient_accumulation_steps: Number of gradient accumulation steps.
  • Resolution Configuration
    • --height: Height of the image/video. Leave empty to enable dynamic resolution.
    • --width: Width of the image/video. Leave empty to enable dynamic resolution.
    • --max_pixels: Maximum pixel area, images larger than this will be scaled down during dynamic resolution.
    • --num_frames: Number of frames for video (video generation models only).
• ACE-Step Specific Parameters
  • --tokenizer_path: Tokenizer path, in format model_id:origin_pattern.
  • --silence_latent_path: Silence latent path, in format model_id:origin_pattern.
  • --initialize_model_on_cpu: Whether to initialize models on CPU.

Example Dataset Copy Section

modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --local_dir ./data/diffsynth_example_dataset

We provide recommended training scripts for each model, please refer to the table in "Model Overview" above. For guidance on writing model training scripts, see Model Training; for more advanced training algorithms, see Training Framework Overview.

En Model Details Anima (`docs/en/Model_Details/Anima.md`) Copy Section

Anima

Anima is an image generation model trained and open-sourced by CircleStone Labs and Comfy Org.

Installation Copy Chapter

Before using this project for model inference and training, please install DiffSynth-Studio first.

git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
pip install -e .

For more installation information, please refer to Install Dependencies.

Quick Start Copy Chapter

The following code demonstrates how to quickly load the circlestone-labs/Anima model for inference. VRAM management is enabled by default, allowing the framework to automatically control model parameter loading based on available VRAM. Minimum 8GB VRAM required.

/* Detailed source-code truncated for AI context efficiency. */

Model Overview Copy Chapter

Special training scripts:

• Differential LoRA Training: doc
• FP8 Precision Training: doc
• Two-Stage Split Training: doc
• End-to-End Direct Distillation: doc

Model Inference Copy Chapter

Models are loaded through AnimaImagePipeline.from_pretrained, see Model Inference for details.

Input parameters for AnimaImagePipeline inference include:

• prompt: Text description of the desired image content.
• negative_prompt: Content to exclude from the generated image (default: "").
• cfg_scale: Classifier-free guidance parameter (default: 4.0).
• input_image: Input image for image-to-image generation (default: None).
• denoising_strength: Controls similarity to input image (default: 1.0).
• height: Image height (must be multiple of 16, default: 1024).
• width: Image width (must be multiple of 16, default: 1024).
• seed: Random seed (default: None).
• rand_device: Device for random noise generation (default: "cpu").
• num_inference_steps: Inference steps (default: 30).
• sigma_shift: Scheduler sigma offset (default: None).
• progress_bar_cmd: Progress bar implementation (default: tqdm.tqdm).

For VRAM constraints, enable VRAM Management. Recommended low-VRAM configurations are provided in the "Model Overview" table above.

Model Training Copy Chapter

Anima models are trained through examples/anima/model_training/train.py with parameters including:

• General Training Parameters
  • Dataset Configuration
    • --dataset_base_path: Dataset root directory.
    • --dataset_metadata_path: Metadata file path.
    • --dataset_repeat: Dataset repetition per epoch.
    • --dataset_num_workers: Dataloader worker count.
    • --data_file_keys: Metadata fields to load (comma-separated).
  • Model Loading
    • --model_paths: Model paths (JSON format).
    • --model_id_with_origin_paths: Model IDs with origin paths (e.g., "anima-team/anima-1B:text_encoder/*.safetensors").
    • --extra_inputs: Additional pipeline inputs (e.g., controlnet_inputs for ControlNet).
    • --fp8_models: FP8-formatted models (same format as --model_paths).
    • --quant_options: Dynamically quantize loaded models. Semicolon-separated entries, each <model_string>:<method>[/<exclude_modules>], where <model_string> matches an entry in --model_paths/--model_id_with_origin_paths, method is a registered method (e.g. bitsandbytes_nf4), and exclude_modules optionally lists layers kept in full precision.
  • Training Configuration
    • --learning_rate: Learning rate.
    • --num_epochs: Training epochs.
    • --trainable_models: Trainable components (e.g., dit, vae, text_encoder).
    • --find_unused_parameters: Handle unused parameters in DDP training.
    • --weight_decay: Weight decay value.
    • --task: Training task (default: sft).
  • Output Configuration
    • --output_path: Model output directory.
    • --remove_prefix_in_ckpt: Remove state dict prefixes.
    • --save_steps: Model saving interval.
  • LoRA Configuration
    • --lora_base_model: Target model for LoRA.
    • --lora_target_modules: Target modules for LoRA.
    • --lora_rank: LoRA rank.
    • --lora_checkpoint: LoRA checkpoint path.
    • --preset_lora_path: Preloaded LoRA checkpoint path.
    • --preset_lora_model: Model to merge LoRA with (e.g., dit).
  • Gradient Configuration
    • --use_gradient_checkpointing: Enable gradient checkpointing.
    • --use_gradient_checkpointing_offload: Offload checkpointing to CPU.
    • --gradient_accumulation_steps: Gradient accumulation steps.
  • Image Resolution
    • --height: Image height (empty for dynamic resolution).
    • --width: Image width (empty for dynamic resolution).
    • --max_pixels: Maximum pixel area for dynamic resolution.
• Anima-Specific Parameters
  • --tokenizer_path: Tokenizer path for text-to-image models.
  • --tokenizer_t5xxl_path: T5-XXL tokenizer path.