FEATURED · 精选文章

DiffSynth-Studio Diffusion Templates 架构详解:基于模板模型的可控生成框架原理与实战

发布时间 / 2026/9/15 12:25:17
来源 / 创域科博编辑部
栏目 / 资讯中心
DiffSynth-Studio Diffusion Templates 架构详解:基于模板模型的可控生成框架原理与实战 DiffSynth-Studio Diffusion Templates 架构详解基于模板模型的可控生成框架原理与实战【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio导读Diffusion Templates 是 DiffSynth-Studio 中面向扩散模型的可控生成插件框架它通过在基础 Diffusion Pipeline 之外引入独立的 Template Pipeline让若干个轻量模板模型以插件形式接管基础模型的部分输入参数从而实现亮度调节、结构控制、图像编辑、局部重绘、超分辨率等可控生成能力。本文以 Understanding_Diffusion_Templates.md 为骨架结合仓库源码template.py、flux2_image.py 等与推理/训练示例深入讲解框架的模块设计、模型能力媒介KV-Cache / Residual / LoRA、Template 模型文件格式以及从推理到训练、上传的完整实战路径。读完本文你将能够理解该框架的架构原理并掌握在 FLUX.2 基础模型上加载、组合、训练 Template 模型的具体方法。框架结构总览Diffusion Templates 框架由 Template Input、Template Model、Template Cache、Template Pipeline 四类模块组成整体结构如下图所示框架包含以下模块设计Template InputTemplate 模型的输入。其格式为 Python 字典其中的字段由每个 Template 模型自身决定例如{scale: 0.8}Template ModelTemplate 模型可从魔搭模型库加载ModelConfig(model_idxxx/xxx)或从本地路径加载ModelConfig(pathxxx)Template CacheTemplate 模型的输出。其格式为 Python 字典其中的字段仅支持对应基础模型 Pipeline 中的输入参数字段Template Pipeline用于调度多个 Template 模型的模块。该模块负责加载 Template 模型、整合多个 Template 模型的输出从源码看TemplatePipeline的实现位于 diffsynth/diffusion/template.pyTemplatePipeline.__init__接收torch_dtype、device、model_configs、lazy_loading四个参数将多个 Template 模型组织为torch.nn.ModuleList其from_pretrained静态方法L154-161是标准的构建入口与基础模型 Pipeline 的加载方式保持一致。框架启用前后的差异未启用 Diffusion Templates 框架时基础模型组件包括 Text Encoder、DiT、VAE 等被加载到 Diffusion Pipeline 中输入 Model Input包括 prompt、height、width 等输出 Model Output例如图像。启用 Diffusion Templates 框架后若干个 Template 模型被加载到 Template Pipeline 中Template Pipeline 输出 Template CacheDiffusion Pipeline 输入参数的子集并交由 Diffusion Pipeline 进行后续的进一步处理。Template Pipeline 通过接管一部分 Diffusion Pipeline 的输入参数来实现可控生成。这一设计在TemplatePipeline.__call__template.py中体现得十分直接框架先调用call_single_side分别计算正向与负向的 Template Cache然后通过inspect.signature(pipe.__call__).parameters检查基础 Pipeline 的输入签名只有属于基础 Pipeline 输入参数列表的字段才会被写入kwargs注入基础 Pipeline其余字段会被打印告警并忽略负向 Cache 则以negative_前缀的字段注入例如negative_kv_cache。这从机制上保证了Template Cache 只能是 Diffusion Pipeline 输入参数的子集这一约束。模型能力媒介为什么 KV-Cache 是最佳选择注意到 Template Cache 的格式被定义为 Diffusion Pipeline 输入参数的子集这是框架通用性设计的基本保证——限制 Template 模型的输入只能是 Diffusion Pipeline 的输入参数。因此需要为 Diffusion Pipeline 设计额外的输入参数作为模型能力媒介。其中KV-Cache 是非常适合 Diffusion 的模型能力媒介理由如下技术路线已经在 LLM Skills 上得到验证LLM 中输入的提示词也会被潜在地转化为 KV-Cache高权限KV-Cache 具有 Diffusion 模型的高权限在生图模型上能够直接影响甚至完全控制生图结果这保证 Diffusion Template 模型具备足够高的能力上限可拼接KV-Cache 可以直接在序列层面拼接让多个 Template 模型同时生效开发成本低KV-Cache 在框架层面的开发量少增加一个 Pipeline 的输入参数并穿透到模型内部即可可以快速适配新的 Diffusion 基础模型。源码中的 KV-Cache 合并与穿透多模型拼接的实现在 template.py 的merge_kv_cache中框架按 KV-Cache 的键名对应 DiT 的 block 名汇总所有 Template 模型产出的 K、V 张量并在dim1序列维度上做torch.concat因此多个 Template 模型的 KV-Cache 可以同时生效且互不干扰。在合并多个 Template 模型输出时merge_template_cachetemplate.py针对不同字段类型采用不同的合并策略kv_cache调用merge_kv_cache做序列拼接lora调用 utils/lora/merge.py 中的merge_lora做 LoRA 权重融合text_embedding调用merge_sequential_embeddings做序列拼接其余字段若只出现在一个 Template 模型输出中则直接透传若多个模型输出同名非媒介字段则打印Conflict detected告警并仅保留第一个。其他可用的模型能力媒介除 KV-Cache 外框架还支持以下媒介Residual残差在 ControlNet 中使用较多适合做点对点的控制与 KV-Cache 相比缺点是不能支持任意分辨率且多个 Residual 融合时可能冲突。LoRA不要把它当成模型的一部分而是把它当成模型的输入参数。LoRA 本质上是一系列张量也可以作为模型能力的媒介。当前支持范围目前仅在 FLUX.2 的 Pipeline 上提供了 KV-Cache 和 LoRA 作为 Template Cache 的支持后续会考虑支持更多模型和更多模型能力媒介。FLUX.2 Pipeline 侧的媒介落地在 diffsynth/pipelines/flux2_image.py 中Flux2ImagePipeline的调用签名新增了对应输入参数kv_cache、negative_kv_cache、lora、negative_lora、extra_text_embedding、negative_extra_text_embedding分别被映射为 DiT 的kv_cache、extra_text_embedding、positive_only_lora、negative_only_lora。在模型内部diffsynth/models/flux2_dit.py 中 KV-Cache 以key torch.concat([key, kv_cache[0]], dim1)的方式拼接到注意力计算的 K/V 上并分别在double_{i}与single_{i}两类 blockflux2_dit.py 与 flux2_dit.py上生效extra_text_embedding则通过构造额外的文本 token 索引与prompt_embeds序列拼接flux2_image.py实现文本层面的额外输入。Template 模型格式一个 Template 模型的格式为Template_Model ├── model.py └── model.safetensors其中model.py是模型的入口model.safetensors是 Template 模型的权重文件。从源码看load_template_modeltemplate.py通过importlib.util.spec_from_file_location动态加载目录下的model.py模块并读取模块级变量若定义了TEMPLATE_MODEL_PATH权重文件相对路径则用load_model加载预训练权重若未定义则直接实例化TEMPLATE_MODEL()得到随机初始化的模型或非模型模块。加载完成后check_template_model_formattemplate.py会强制校验模型必须实现process_inputs与forward且两者都必须包含**kwargs参数否则抛出NotImplementedError。process_inputs与forward构成完整的 Template 模型推理过程——Template Input 先经过process_inputs无梯度计算再进入forward得到 Template Cache这种拆分是为了在训练中更容易适配两阶段拆分训练。关于如何构建 Template 模型请参考文档 Template 模型训练已发布模型的推理与训练入口可参考 Introducing_Diffusion_Templates.md 中的模型清单结构控制、亮度调节、色彩调节、图像编辑、超分辨率、锐利激发、美学对齐、局部重绘、内容参考、年龄控制等。推理实战在 FLUX.2 上启用 Template 模型下面以基础模型 black-forest-labs/FLUX.2-klein-base-4Bfrom diffsynth.diffusion.template import TemplatePipeline from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig import torch # 1. 加载基础模型 Pipeline pipe Flux2ImagePipeline.from_pretrained( torch_dtypetorch.bfloat16, devicecuda, model_configs[ ModelConfig(model_idblack-forest-labs/FLUX.2-klein-base-4B, origin_file_patterntransformer/*.safetensors), ModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patterntext_encoder/*.safetensors), ModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patternvae/diffusion_pytorch_model.safetensors), ], tokenizer_configModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patterntokenizer/), ) # 2. 加载 Template Pipeline可从魔搭加载或本地路径加载 template TemplatePipeline.from_pretrained( torch_dtypetorch.bfloat16, devicecuda, model_configs[ModelConfig(model_idDiffSynth-Studio/Template-KleinBase4B-Brightness)], ) # 3. 推理把 pipe 的输入参数转移到 template 中并添加 template_inputs image template( pipe, promptA cat is sitting on a stone., seed0, cfg_scale4, num_inference_steps50, template_inputs[{scale: 0.7}], # scale 提高亮度 negative_template_inputs[{scale: 0.5}], # CFG 增强 ) image.save(image_Brightness_light.jpg)注意需将pipe的输入参数转移到template_pipeline中并添加template_inputs。template_inputs中的字段如scale由每个 Template 模型自身定义会作为 Template Input 传入对应模型的process_inputs。Template 模型的 CFG 增强Template 模型可以开启 CFGClassifier-Free Guidance使其控制效果更明显。在TemplatePipeline的输入参数中添加negative_template_inputs模型就会对比正向与负向两侧的差异生成控制效果更明显的图像实现细节见 template.py 中call_single_side分别计算正负两侧 Cache 的逻辑image template( pipe, promptA cat is sitting on a stone., seed0, cfg_scale4, num_inference_steps50, template_inputs[{scale: 0.8}], negative_template_inputs[{scale: 0.5}], )各能力类型的 Template 模型用法仓库中提供了多种能力的 Template 推理示例见 examples/flux2/model_inference/ 目录低显存版本见 examples/flux2/model_inference_low_vram/其template_inputs结构各不相同| 能力 | 示例文件 | template_inputs 关键字段 | | - | - | - | | 亮度调节 | Template-KleinBase4B-Brightness.py |{scale: 0.7}0~1 区间越大越亮 | | 结构控制 | Template-KleinBase4B-ControlNet.py |{image: PIL.Image, prompt: str}深度图提示词 | | 美学对齐 | Template-KleinBase4B-Aesthetic.py |{lora_ids: [...], lora_scales: 1.0, merge_type: mean}LoRA 媒介 | | 图像编辑 | Template-KleinBase4B-Edit.py |{image: PIL.Image, prompt: str}参考图编辑指令 | | 局部重绘 | Template-KleinBase4B-Inpaint.py |{image: PIL.Image, mask: PIL.Image, force_inpaint: True}|其中美学对齐LoRA 媒介示例有一个关键前置步骤必须先调用pipe.dit pipe.enable_lora_hot_loading(pipe.dit)开启 LoRA 热加载示例代码中标注了# Important!否则会导致 LoRA 权重在多次调用间叠加。ControlNet、Edit、Inpaint 等示例还需通过dataset_snapshot_download(DiffSynth-Studio/examples_in_diffsynth, allow_file_pattern[templates/*], ...)下载示例输入图片深度图、参考图、蒙版到data/examples/templates/。启用多个 Template 模型TemplatePipeline可以加载多个 Template 模型通过model_configs列表推理时在template_inputs中使用model_id字段区分每个 Template 模型的输入call_single_sidetemplate.py会按model_id调度并缓存对应的模型最后统一合并 Template Cache。多个 Template 模型可以叠加生效——例如超分辨率 锐利激发、结构控制 美学对齐 锐利激发的组合应用可参考 Introducing_Diffusion_Templates.md 中的效果一览。低显存支持Template 模型暂不支持主框架的显存管理——check_vram_configtemplate.py会对offload_device、computation_dtype等 VRAM 配置参数发出TemplatePipeline doesnt support VRAM management警告并忽略。但可以使用惰性加载添加参数lazy_loadingTrue仅在需要推理时加载对应的 Template 模型。这在启用多个 Template 模型时可以显著降低显存需求显存占用峰值为单个 Template 模型的显存占用量template TemplatePipeline.from_pretrained( torch_dtypetorch.bfloat16, devicecuda, model_configs[ModelConfig(model_idDiffSynth-Studio/Template-KleinBase4B-Brightness)], lazy_loadingTrue, )基础模型的 Pipeline 与 Template Pipeline完全独立可按需对基础模型开启显存管理参考 VRAM_management.md。构建并训练新的 Template 模型Template 模型组件格式一个 Template 模型与一个模型库或一个本地文件夹绑定模型库中有代码文件model.py作为唯一入口。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推理时的数据流Template Input → process_inputs → forward → Template Cache与训练时的数据流Dataset → TEMPLATE_DATA_PROCESSOR → process_inputs → forward → Template Cache分别如下TEMPLATE_MODELTEMPLATE_MODEL是 Template 模型的代码实现需继承torch.nn.Module并编写process_inputs与forward两个函数process_inputs需带有装饰器torch.no_grad()进行不包含梯度的计算forward需包含训练模型所需的全部梯度计算过程其输入与process_inputs的输出相同。两者都必须包含**kwargs保证兼容性这也是check_template_model_format的强制校验项并预留以下参数pipe如需在process_inputs与forward中和基础模型 Pipeline 进行交互例如调用基础模型 Pipeline 中的文本编码器进行编码可在输入参数中增加字段pipeuse_gradient_checkpointing/use_gradient_checkpointing_offload如需在训练中启用 Gradient Checkpointing可在forward的输入参数中增加这两个字段多个 Template 模型需通过model_id区分 Template Inputs请不要在process_inputs与forward的输入参数中使用这个字段。TEMPLATE_MODEL_PATH可选项TEMPLATE_MODEL_PATH是模型预训练权重文件的相对路径TEMPLATE_MODEL_PATH model.safetensors如需从多个模型文件中加载可使用列表TEMPLATE_MODEL_PATH [ model-00001-of-00003.safetensors, model-00002-of-00003.safetensors, model-00003-of-00003.safetensors, ]如果需要随机初始化模型参数模型还未训练或不需要初始化模型参数可将其设置为None或不设置TEMPLATE_MODEL_PATH NoneTEMPLATE_DATA_PROCESSOR可选项如需使用 DiffSynth-Studio 训练 Template 模型则需构建训练数据集数据集中的metadata.json或metadata.jsonl包含template_inputs字段。该字段并不是直接输入给 Template 模型process_inputs的参数而是提供给TEMPLATE_DATA_PROCESSOR的输入参数由TEMPLATE_DATA_PROCESSOR计算出输入给 Template 模型process_inputs的参数。例如亮度控制模型DiffSynth-Studio/Template-KleinBase4B-Brightness的输入参数是scale图像的亮度数值可直接写在metadata.json中此时TEMPLATE_DATA_PROCESSOR只需透传参数[ { 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也可在metadata.json中填写图像路径在训练过程中直接计算scale例如取图像像素均值归一化见 examples/flux2/model_training/scripts/brightness/model.py 的DataAnnotator[ { 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训练 Template 模型Template 模型可训练的充分条件是Template Cache 中的变量计算与基础模型 Pipeline 完全解耦——这些变量在推理过程中输入给基础模型 Pipeline 后不会参与任何 Pipeline Unit 的计算直达model_fn。以基础模型 FLUX.2-klein-base-4B 为例训练脚本完整样例见 Template-KleinBase4B-Brightness.sh训练入口为 examples/flux2/model_training/train.py中的关键参数--extra_inputs额外输入。训练文生图模型的 Template 模型时只需填template_inputs训练图像编辑模型的 Template 模型时需填edit_image,template_inputs--template_model_id_or_pathTemplate 模型的魔搭模型 ID 或本地路径。框架会优先匹配本地路径若本地路径不存在则从魔搭模型库中下载该模型填写模型 ID 时以 : 结尾例如DiffSynth-Studio/Template-KleinBase4B-Brightness:--remove_prefix_in_ckpt保存模型文件时移除的 state dict 变量名前缀填pipe.template_model.即可--trainable_models可训练模型填写template_model即可若只需训练其中的某个组件则需填写template_model.xxx,template_model.yyy以逗号分隔以下是一个样例训练脚本它会自动下载一个样例数据集随机初始化模型权重后开始训练亮度控制模型modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include flux2/Template-KleinBase4B-Brightness/* --local_dir ./data/diffsynth_example_dataset accelerate launch examples/flux2/model_training/train.py \ --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness \ --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness/metadata.jsonl \ --extra_inputs template_inputs \ --max_pixels 1048576 \ --dataset_repeat 50 \ --model_id_with_origin_paths black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors \ --template_model_id_or_path 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以亮度控制模型为例其模型实现 examples/flux2/model_training/scripts/brightness/model.py 清晰地展示了 KV-Cache 媒介的生成方式ValueFormatModel为 DiT 的每个double_{i}5 个与single_{i}20 个block 分别注册一组SingleValueEncoderproj_k/proj_vforward中把标量value通过正弦位置编码timestep embedding MLP 映射为定长序列的 K/V 张量最终产出kv_cache[block_name] (k, v)字典。这正是框架架构中Template 模型输出字段由自身决定、但最终注入基础 Pipeline 输入参数的典型实现。与基础模型 Pipeline 组件交互Diffusion Template 框架允许 Template 模型与基础模型 Pipeline 进行交互。例如你可能需要使用基础模型 Pipeline 中的 text encoder 对文本进行编码此时在process_inputs和forward中使用预留字段pipe即可import torch 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(text_emb) 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使用非训练的模型组件在设计 Template 模型时如果需要使用预训练的模型且不希望在训练过程中更新这部分参数例如在__init__中加载一个固定的image_encoder仅训练其后的mlp此时需在训练命令中通过参数--trainable_models template_model.mlp设置为仅训练mlp部分import torch 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在低显存的设备上训练框架支持将 Template 模型的训练拆分为两阶段第一阶段进行无梯度计算第二阶段进行梯度更新更多信息请参考两阶段拆分训练以下是样例脚本modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include flux2/Template-KleinBase4B-Brightness/* --local_dir ./data/diffsynth_example_dataset # 阶段一数据预处理无梯度计算输出 Cache accelerate launch examples/flux2/model_training/train.py \ --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness \ --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Brightness/metadata.jsonl \ --extra_inputs template_inputs \ --max_pixels 1048576 \ --dataset_repeat 1 \ --model_id_with_origin_paths black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors \ --template_model_id_or_path DiffSynth-Studio/Template-KleinBase4B-Brightness: \ --tokenizer_path black-forest-labs/FLUX.2-klein-4B:tokenizer/ \ --learning_rate 1e-4 \ --num_epochs 2 \ --remove_prefix_in_ckpt pipe.template_model. \ --output_path ./models/train/Template-KleinBase4B-Brightness_full_cache \ --trainable_models template_model \ --use_gradient_checkpointing \ --find_unused_parameters \ --task sft:data_process # 阶段二梯度更新仅加载 Transformer 与缓存的数据 accelerate launch examples/flux2/model_training/train.py \ --dataset_base_path ./models/train/Template-KleinBase4B-Brightness_full_cache \ --extra_inputs template_inputs \ --max_pixels 1048576 \ --dataset_repeat 50 \ --model_id_with_origin_paths black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors \ --template_model_id_or_path DiffSynth-Studio/Template-KleinBase4B-Brightness: \ --tokenizer_path black-forest-labs/FLUX.2-klein-4B:tokenizer/ \ --learning_rate 1e-4 \ --num_epochs 2 \ --remove_prefix_in_ckpt pipe.template_model. \ --output_path ./models/train/Template-KleinBase4B-Brightness_full \ --trainable_models template_model \ --use_gradient_checkpointing \ --find_unused_parameters \ --task sft:train两阶段拆分训练可以降低显存需求、提高训练速度训练过程是无损精度的但需要较大硬盘空间用于存储 Cache 文件。如需进一步减少显存需求可开启 fp8 精度在两阶段训练中添加参数--fp8_models black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors和--fp8_models black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors即可。fp8 精度只能在非训练模型组件上启用且存在少量误差。上传 Template 模型完成训练后按照以下步骤可上传 Template 模型到魔搭社区供更多人下载使用。Step 1在model.py中填入训练好的模型文件名TEMPLATE_MODEL_PATH model.safetensorsStep 2上传model.py--token ms-xxx在魔搭社区个人访问令牌页面获取modelscope upload user_name/your_model_id /path/to/your/model.py model.py --token ms-xxxStep 3确认模型文件例如epoch-1.safetensors、step-2000.safetensors。注意DiffSynth-Studio 保存的模型文件中只包含可训练的参数如果模型中包括非训练参数则需要重新将非训练的模型参数打包才能进行推理可以通过以下代码打包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_dtypetorch.bfloat16, devicecpu) state_dict load_state_dict(path/to/your/ckpt/epoch-1.safetensors, torch_dtypetorch.bfloat16, devicecpu) state_dict.update(model.state_dict()) save_file(state_dict, model.safetensors)Step 4上传模型文件modelscope upload user_name/your_model_id /path/to/your/model/epoch-1.safetensors model.safetensors --token ms-xxxStep 5验证模型推理效果from diffsynth.diffusion.template import TemplatePipeline from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig import torch pipe Flux2ImagePipeline.from_pretrained( torch_dtypetorch.bfloat16, devicecuda, model_configs[ ModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patterntext_encoder/*.safetensors), ModelConfig(model_idblack-forest-labs/FLUX.2-klein-base-4B, origin_file_patterntransformer/*.safetensors), ModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patternvae/diffusion_pytorch_model.safetensors), ], tokenizer_configModelConfig(model_idblack-forest-labs/FLUX.2-klein-4B, origin_file_patterntokenizer/), ) template_pipeline TemplatePipeline.from_pretrained( torch_dtypetorch.bfloat16, devicecuda, model_configs[ModelConfig(model_iduser_name/your_model_id)], ) image template_pipeline( pipe, prompta cat, seed0, cfg_scale4, height1024, width1024, template_inputs[{xxx}], ) image.save(image.png)总结Diffusion Templates 框架通过Template Pipeline 接管 Diffusion Pipeline 部分输入参数的设计实现了通用、可组合、可训练的可控生成能力以 KV-Cache 为核心模型能力媒介辅以 Residual、LoRA让多个 Template 模型可以在序列层面拼接并同时生效以model.py model.safetensors的极简模型格式与process_inputs → forward的两段式推理接口兼顾了推理的灵活性与训练含两阶段拆分训练的可扩展性。目前框架已在 FLUX.2 基础模型上落地了亮度、色彩、结构控制、编辑、重绘、超分辨率、美学对齐、内容参考、年龄控制等十余种能力其架构与实现可进一步参考 Introducing_Diffusion_Templates.md、Template_Model_Inference.md 与 Template_Model_Training.md。【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻