JoyImageEditPlusTransformer3DModel 深度解析Diffusers 多图像指令编辑的 3D 扩散 Transformer【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers导读本文围绕 transformer_joyimage_edit_plus.md 中定义的JoyImageEditPlusTransformer3DModel系统讲解它的加载方式、双流 Transformer 架构、6D Patch 化输入协议、三维旋转位置编码RoPE与配置参数并结合 pipeline_joyimage_edit_plus.py 与 test_models_transformer_joyimage_edit_plus.py 等源码与测试剖析它在多参考图像编辑任务中的实际调用链路。读完本文你将掌握该模型的结构原理、每个配置项的含义与取值以及如何在diffusers中独立加载并驱动它完成多图指令编辑。模型定位为多图像编辑设计的 3D 扩散 TransformerJoyImageEditPlusTransformer3DModel是diffusers中服务于 JoyAI-Image-Edit-Plus 多图像指令编辑任务的 Transformer 主干注册在 src/diffusers/models/init.py 的diffusers.models.transformers.transformer_joyimage_edit_plus模块下并可在顶层通过from diffusers import JoyImageEditPlusTransformer3DModel直接导入。它的核心任务是接受多张参考图像不同分辨率与目标噪声在潜在空间中执行联合去噪从而生成一张按照文本指令组合参考图像元素的新图像。与常见的单图像编辑模型不同该模型一次可处理 15 张参考图并且允许每张参考图具有不同的分辨率——这得益于其独特的独立 Patch 化 拼接设计详见下文6D Patch 化输入一节。从 pipeline_joyimage_edit_plus.py 的类注释可以确认它的完整工作语境调度器FlowMatchEulerDiscreteScheduler流匹配 Euler 离散调度器VAEAutoencoderKLWan将像素图像编解码为潜在表示文本编码器Qwen3VLForConditionalGeneration多模态文本编码支持行内图像理解Transformer即本文主角JoyImageEditPlusTransformer3DModelMMDiT 架构负责去噪处理器Qwen3VLProcessor处理文本 图像的多模态输入。from diffusers import JoyImageEditPlusTransformer3DModel transformer JoyImageEditPlusTransformer3DModel.from_pretrained( jdopensource/JoyAI-Image-Edit-Plus-Diffusers, subfoldertransformer, dtypetorch.bfloat16 )上述代码即关联文档给出的官方加载方式从jdopensource/JoyAI-Image-Edit-Plus-Diffusers仓库的transformer子目录加载预训练权重并以bfloat16精度载入。from_pretrained是ModelMixin提供的标准入口源码见 transformer_joyimage_edit_plus.py它会读取仓库中的config.json恢复完整配置并加载权重。架构解析双流Double-Stream联合注意力设计JoyImageEditPlusTransformer3DModel继承ModelMixin、ConfigMixin与AttentionMixin见 transformer_joyimage_edit_plus.py整体是一个20 层双流 Transformer默认num_layers20。它的网络结构自上而下分为五层1. 输入投影3D 卷积 Patch 化img_inimg_in是一个nn.Conv3d(in_channels, hidden_size, kernel_sizepatch_size, stridepatch_size)L387即用卷积同时完成 Patch 切分与通道投影。默认patch_size[1, 2, 2]表示沿时间维切 1 个、空间高宽各切 2 个将潜在张量变成隐藏向量序列。2. 条件嵌入器时间 文本联合嵌入condition_embedderJoyImageEditPlusTimeTextImageEmbeddingL282-L314完成三件事时间步通过Timesteps(num_channels256, flip_sin_to_cosTrue)做正弦余弦编码再经TimestepEmbedding映射到hidden_sizeSiLU激活后经time_proj线性层投影为hidden_size * 6维的调制向量vec之后按 6 份切分作为每个 Transformer 块的调制条件文本嵌入经PixArtAlphaTextProjection(text_embed_dim, dim, act_fngelu_tanh)投影到hidden_size维度。输出三元组(temb, timestep_proj, encoder_hidden_states)其中temb与timestep_proj目前分别服务于不同用途调制向量的中间表示来自timestep_proj。3. 双流 Transformer 块图像流与文本流并行每个JoyImageEditPlusTransformerBlockL194-L279同时维护图像流hidden_states与文本流encoder_hidden_states两套并行的调制、归一化与 MLP 分支每条流各有一套JoyImageEditPlusModulatefactor6调制器产出shift1/scale1/gate1与shift2/scale2/gate2共 6 个调制向量采用 Wan 风格的可学习调制表modulate_table初始化为zeros(1, factor, hidden_size) / hidden_size**0.5前向时将条件信号加到表上再按 factor 切分L56-L74归一化使用FP32LayerNormelementwise_affineFalse保证在 bfloat16 训练/推理下 LayerNorm 仍以 FP32 精度计算数值更稳定MLP 使用FeedForward激活函数为gelu-approximatemlp_width_ratio4.0时隐藏维度为hidden_size * 4 12288。每条流的调制公式为modulated norm(x) * (1 scale) shift再经联合注意力后按gate缩放做残差相加随后进入第二条流的调制 MLP 残差。4. 联合注意力图像与文本的 QKV 拼接核心JoyImageEditPlusAttentionL145-L191是本模型区别于普通文本条件扩散模型的关键。它对图像流与文本流分别计算 QKV图像侧img_attn_qkv线性层 img_attn_q_norm/img_attn_k_normRMSNorm 归一化img_attn_proj输出投影文本侧txt_attn_qkv线性层 txt_attn_q_norm/txt_attn_k_normtxt_attn_proj输出投影。在JoyImageEditPlusAttnProcessorL77-L142中图像与文本的 Q/K/V 各自归一化图像侧还可选施加 RoPE后沿序列维拼接成联合 Q/K/V调用dispatch_attention_fn执行联合注意力最后把输出按序列位置切分回图像部分与文本部分分别投影。这种双流各自 QKV 联合注意力的结构即 MMDiT 风格的文本-图像信息交互方式。需要指出的是该注意力处理器强制要求提供encoder_hidden_states否则直接抛出ValueErrorL91-L92——这从源码层面印证了该模型是条件模型文本嵌入是必填输入而非可选项。5. 输出层归一化 线性投影回 Patch最后经过norm_outFP32LayerNorm与proj_out nn.Linear(hidden_size, out_channels * prod(patch_size))L408-L409把隐藏状态投影回out_channels默认等于in_channels16乘以 Patch 体积的维度再重塑为 6D Patch 张量封装为Transformer2DModelOutput(sampleimg)返回L531-L539。Transformer2DModelOutput定义于 models/modeling_outputs.py[[autodoc]]指向models.modeling_outputs.Transformer2DModelOutput。6D Patch 化输入协议多图变分辨率的实现基础关联文档明确指出该模型的输入格式为[B, max_patches, C, pt, ph, pw]的 6D 填充 Patch 张量见类 docstringL321-L324。理解这一协议是理解整个模型的关键Bbatch 大小max_patchesbatch 内最大的 Patch 序列长度各样本按此长度对齐填充C输入通道数默认 16pt, ph, pw每个 Patch 的尺寸默认 1×2×2。目标噪声与每张参考图被独立 Patch 化然后沿序列维拼接成一条扁平的 Patch 序列。由于每张参考图可以有不同的分辨率各自的 Patch 数量也不同因此需要对 batch 内所有样本统一填充到max_patches。这种先独立 Patch、再拼接、再填充的做法正是模型能够支持变分辨率参考图的原因。这一协议在 pipeline 的prepare_latentspipeline_joyimage_edit_plus.py L274-L374中有完整的工程实现目标位置采样随机噪声形状(C, 1, H, W)参考图经JoyImageEditImageProcessor预处理后由 VAE 编码为潜在表示并按latents_mean/latents_std归一化每个组件目标噪声 各参考图按patch_size切分为(l_t, l_h, l_w)网格重塑为 Patch 序列同一样本的所有组件 Patch 拼接并记录每个组件的(t, h, w)元组到shape_listbatch 内按max_patches零填充对齐同时生成target_mask标记目标 Patch 的位置。shape_list是前向传播的必传参数shape_list: list[list[tuple[int, int, int]]]它按样本记录每个组件目标 各参考图的 Patch 网格尺寸用于构建 RoPE 与注意力掩码。前向传播全流程六步拆解结合 forward 的实现模型前向可拆为六步条件嵌入condition_embedder从timestep与encoder_hidden_states计算调制向量vecunflatten 为 6 组与投影后的文本嵌入txtPatch 化6D 输入先reshape(batch_size * max_num_patches, ...)经img_inConv3d卷积后 reshape 为(B, max_patches, D)逐组件 RoPE对每个样本的每个组件按其(t, h, w)调用_get_rotary_pos_embed_for_range生成 3D 旋转位置编码。关键细节组件间的 RoPE 时间偏移是累计的——current_t_offset从 0 开始每处理完一个组件就累加其时间维 Patch 数L481-L492从而在位置编码层面区分目标噪声与第几张参考图不足max_patches的部分用 cos1、sin0 填充等价于不旋转注意力掩码由encoder_hidden_states_mask与图像 Patch 掩码拼接成[B, 1, 1, img_seq txt_seq]的布尔掩码L508-L516双层循环依次运行 20 个JoyImageEditPlusTransformerBlock每个块内部执行调制 → 联合注意力 → 门控残差 → FFN 的流程开启gradient_checkpointing时通过_gradient_checkpointing_func计算图重放以节省显存输出投影proj_out(norm_out(img))后 reshape 回 6D Patch 格式返回Transformer2DModelOutput。关于 RoPE 的实现细节_get_rotary_pos_embed_for_rangeL417-L442按rope_dim_list默认[16, 56, 56]分别对应 t/h/w 三个维度的频率分量数生成网格频率基数为theta256频率公式为1 / theta^(arange(0, dim, 2) / dim)即标准的 NTK 风格旋转位置编码_apply_rotary_emb_batchedL35-L53以批量方式支持[B, S, D]频率张量对 Q/K 施加旋转这也是为多图输入设计的批量化 RoPE 实现。配置参数详解JoyImageEditPlusTransformer3DModel.__init__通过register_to_config注册全部配置L362-L376保存于config.json。默认值面向完整模型hidden_size3072、24 头、20 层参数默认值说明patch_size[1, 2, 2]潜在输入的 Patch 大小沿(t, h, w)三维in_channels16输入潜在张量的通道数与 VAE 潜在维度一致out_channelsNone输出通道数不指定时默认等于in_channelshidden_size3072隐藏表示维度num_attention_heads24注意力头数text_dim4096文本编码器输出的特征维度mlp_width_ratio4.0MLP 隐藏维度相对hidden_size的比例num_layers20双流 Transformer 块数量rope_dim_list[16, 56, 56]3D 旋转位置编码在(t, h, w)上的频率分量维度rope_typerope旋转位置编码类型theta256旋转位置编码的基频源码层面的约束与校验hidden_size必须能被num_attention_heads整除否则抛出ValueErrorL381-L385默认 3072 / 24 128 维每头rope_dim_list为None时自动回退为[head_dim // 3] * 3L425-L426若手动指定rope_dim_list各维度的和应覆盖注意力头维度测试中即用小值[4, 6, 6]对应头维 16。此外模型类还声明了一批与 diffusers 高级特性联动的类属性L351-L360_skip_layerwise_casting_patterns [img_in, condition_embedder, norm]这些层在 layerwise casting 时跳过精度转换_keep_in_fp32_modules [time_embedder, norm1, norm2, norm_out]时间嵌入与各归一化层保持 FP32_supports_gradient_checkpointing True支持梯度检查点_no_split_modules [JoyImageEditPlusTransformerBlock]设备切分/offload 时以 Transformer 块为最小单元_repeated_blocks声明重复块结构。在__init__末尾所有块都被显式设置为JoyImageEditPlusAttnProcessorL413-L415确保批量 RoPE 处理器在所有层统一生效。在 Pipeline 中的实际调用链路模型单独加载之外更常见的用法是作为JoyImageEditPlusPipeline的组成部分被调用。完整的端到端示例见 joyimage_edit_plus.mdimport torch from PIL import Image from diffusers import JoyImageEditPlusPipeline pipeline JoyImageEditPlusPipeline.from_pretrained( jdopensource/JoyAI-Image-Edit-Plus-Diffusers, dtypetorch.bfloat16 ) pipeline.to(cuda) # 或 mps、xpu、cpu images [ Image.open(reference_0.png).convert(RGB), Image.open(reference_1.png).convert(RGB), ] target_h, target_w pipeline.image_processor.get_default_height_width(images[-1]) output pipeline( imagesimages, promptCombine the person from the second image with the scene from the first image., negative_promptlow quality, blurry, deformed, heighttarget_h, widthtarget_w, num_inference_steps30, guidance_scale4.0, generatortorch.Generator(cuda).manual_seed(42), ).images[0] output.save(joyimage_edit_plus_output.png)在去噪循环中pipeline_joyimage_edit_plus.py L666-L724Transformer 的调用方式与单图模型有显著差异参考图保持干净每步去噪前用clean_reference_backup恢复参考图对应的 Patchlatents[~target_mask] clean_reference_backup[~target_mask]L677即只有目标区域参与噪声演化参考区域始终是原始编码CFG 展开启用 classifier-free guidanceguidance_scale 1时模型输入与shape_list均复制一份无条件 有条件时间步同步 repeatCFG 组合带范数重缩放comb_pred noise_pred_uncond guidance_scale * (noise_pred_text - noise_pred_uncond)随后按cond_norm / noise_norm.clamp_min(1e-6)对组合预测做范数重缩放L702-L707这是稳定 CFG 输出尺度的一类常用技巧调度器步进FlowMatchEulerDiscreteScheduler.step更新潜在表示循环结束后仅取出目标 Patch 区域重组潜在张量经 VAE 解码得到最终图像。同时需要注意height与width必须能被 VAE 空间缩放因子整除check_inputs中校验L410-L413。实际分辨率由JoyImageEditImageProcessor的分辨率桶bucket机制决定find_best_bucket在 1024 基准的桶列表覆盖 512×1792 到 2048×512 的多种宽高比中选取与输入图像宽高比最接近的尺寸见 image_processor.py。测试验证模型级与 Pipeline 级双重覆盖仓库为该模型提供了完整的测试体系可用于理解其正确的输入/输出约定模型级测试test_models_transformer_joyimage_edit_plus.pyget_dummy_inputs使用hidden_states[B1, max_patches2, C16, pt1, ph2, pw2]、encoder_hidden_states[B, 12, 16]、timestep与shape_list[[(1,1,1), (1,1,1)]]目标 1 张参考图输出形状为(2, 16, 1, 2, 2)同时继承ModelTesterMixin、MemoryTesterMixin、TrainingTesterMixin验证梯度检查点、AttentionTesterMixin、TorchCompileTesterMixin等通用测试基类Pipeline 级测试test_joyimage_edit_plus.py使用微型随机模型组件huangfeice/tiny-random-Qwen3VLForConditionalGeneration验证多图输入的端到端流程get_dummy_inputs传入两张 32×32 参考图与指令 combine the two images测试还通过 patchfind_best_bucket把分辨率桶固定为 32×32避免 dummy 输入被放大到 1024 级分辨率。这两层测试从独立模型前向与完整管线推理两个维度印证了 6D Patch 协议、shape_list传参和双流注意力的正确性。实践要点与注意事项基于以上源码分析使用JoyImageEditPlusTransformer3DModel时有几点值得注意必传条件前向时必须提供timestep、encoder_hidden_states与shape_list且JoyImageEditPlusAttnProcessor强制要求文本嵌入非空encoder_hidden_states_mask为可选项用于文本 token 的注意力掩码精度策略模型默认以 bfloat16 加载官方示例LayerNorm 与时间嵌入层自动保持 FP32无需手动干预_skip_layerwise_casting_patterns与_keep_in_fp32_modules已为低精度推理做了配置显存优化支持梯度检查点_supports_gradient_checkpointingTrue与模块级设备 offload_no_split_modules在长序列多参考图 长文本场景下可显著降低显存占用pipeline 层面还提供 CPU offloadmodel_cpu_offload_seq text_encoder-transformer-vae分辨率约束输出分辨率需为 VAE 空间缩放因子默认 8的整数倍实际运行时建议沿用 pipeline 的 bucket 机制确定height/width版本适配pipeline 中_get_last_decoder_hidden_states通过注册 forward hook 直接抓取 Qwen3VL 最后一个解码器层的 pre-norm 输出L197-L227以规避 transformers 4.57 与 5.x 之间hidden_states语义变化导致的约 10 倍尺度差异——在自行组装 pipeline 时应留意 transformers 版本与这一适配逻辑的对应关系。若要在不加载完整 pipeline 的情况下独立使用该模型例如自定义训练或特征提取可直接以JoyImageEditPlusTransformer3DModel构造实例并传入 6D Patch 张量参考模型测试中的 dummy 输入构造方式即可快速验证前向通路。关联源码与文档索引模型实现transformer_joyimage_edit_plus.pyPipeline 实现pipeline_joyimage_edit_plus.py图像预处理分辨率桶image_processor.pyPipeline 输出类型pipeline_output.py模型注册src/diffusers/models/init.py模型级测试test_models_transformer_joyimage_edit_plus.pyPipeline 级测试test_joyimage_edit_plus.py端到端使用文档joyimage_edit_plus.md【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考