mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
408ce8db29
* feat: add image edit plus
* refactor: remove debug code
* fix: address review issues for JoyImage Edit Plus
- Remove einops dependency: replace rearrange with reshape/permute
- Remove sglang-specific code from transformer forward
- Remove unused import inspect from transformer
- Fix hardcoded device_type="cuda" to use device.type
- Simplify scheduler sigma math: delegate to retrieve_timesteps
- Remove unused enable_denormalization parameter
- Fix callback latents variable binding
- Fix output_type="pt" to return stacked tensor
- Set return_dict default to True in transformer forward
- Add dummy objects for JoyImageEditPlus classes
- Add transformer and pipeline test files
* fix: add missing newline at end of pipeline_output.py
* fix: add missing newline at end of pipeline_output.py
* doc: add joyimage-edit-plus doc
* refactor: update code format
* refactor: update code formate
* refactor: merge edit-plus conversion script into convert_joyimage_edit_to_diffusers.py
JoyImage Edit and Edit Plus share identical VAE and transformer weight
layouts — only the target model class differs. Consolidate both into a
single script with a --model_type flag (edit | edit_plus) instead of
maintaining two nearly-duplicate files.
* revert: remove unnecessary attention_mask change from transformer_joyimage.py
* Apply style fixes
* fix: remove stale Copied-from annotations that diverged from source
* fix: address CI check failures for edit-plus PR
- Add missing Returns section to JoyImageEditPlusTransformer3DModel.forward docstring
- Fix alphabetical ordering of dummy classes in dummy_pt_objects.py
* fix: use VAE dtype instead of float32 cast in edit-plus pipeline
Replace `torch.autocast(dtype=torch.float32)` + `.float()` with
`.to(self.vae.dtype)` for both VAE encode and decode calls.
The previous approach caused dtype mismatch (float32 input vs bfloat16
bias) on CPU where autocast does not automatically cast conv weights,
breaking the CI `test_layerwise_casting_inference` test.
* refactor: address dg845 review comments on edit-plus
Transformer:
- Simplify RoPE by removing dead txt_freqs path
- Fix patch_size type hint to list[int]
- Use self.config.xxx instead of explicit attributes
Pipeline:
- Merge VaeImageProcessor into JoyImageEditImageProcessor
- Inline normalize/denormalize latents at call sites
- Use image_processor.preprocess() for reference images
- Remove duplicate resize_center_crop in prepare_latents
- Rename padded_latents to latents in __call__
- Support pre-computed latents parameter
- Remove redundant padding token zeroing (attention mask suffices)
- Use image_processor.postprocess() for output conversion
- Return (image,) tuple when return_dict=False
* fix: guard against division-by-zero in CFG norm rescaling
Add .clamp_min(1e-6) to noise_norm to prevent NaN when comb_pred
is all-zeros (e.g. padded patches), matching pipeline_joyimage_edit.py.
* style: run ruff format on edit-plus pipeline
* refactor: address dg845 second-round review comments
- Fix vae_image_processor -> image_processor in docs
- Remove empty __init__ from AttnProcessor
- Restore # Copied from on retrieve_timesteps and sync with source
- Add check_inputs method for basic input validation
- Remove unnecessary test skip for test_from_save_pretrained_dtype_inference
* Apply style fixes
---------
Co-authored-by: tangyanfei.8 <tangyanfei.8@jd.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
356 lines
15 KiB
Python
356 lines
15 KiB
Python
"""Convert JoyImage Edit / Edit Plus checkpoints to diffusers format.
|
|
|
|
Supports both JoyImage-Edit (single-image editing) and JoyImage-Edit-Plus
|
|
(multi-image editing). The transformer weight layout is identical; only the
|
|
target model class and pipeline differ.
|
|
|
|
Usage:
|
|
# Convert JoyImage Edit (default)
|
|
python convert_joyimage_edit_to_diffusers.py \
|
|
--transformer_ckpt_path /path/to/transformer.pt \
|
|
--vae_ckpt_path /path/to/vae.pt \
|
|
--text_encoder_path Qwen/Qwen3-VL-8B-Instruct \
|
|
--output_path /path/to/output \
|
|
--save_pipeline
|
|
|
|
# Convert JoyImage Edit Plus
|
|
python convert_joyimage_edit_to_diffusers.py \
|
|
--model_type edit_plus \
|
|
--transformer_ckpt_path /path/to/transformer.pt \
|
|
--vae_ckpt_path /path/to/vae.pt \
|
|
--text_encoder_path Qwen/Qwen3-VL-8B-Instruct \
|
|
--output_path /path/to/output \
|
|
--save_pipeline
|
|
"""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
from accelerate import init_empty_weights
|
|
from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration
|
|
|
|
from diffusers import (
|
|
AutoencoderKLWan,
|
|
JoyImageEditPipeline,
|
|
JoyImageEditTransformer3DModel,
|
|
)
|
|
from diffusers.models.transformers.transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel
|
|
from diffusers.pipelines.joyimage.pipeline_joyimage_edit_plus import JoyImageEditPlusPipeline
|
|
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
|
|
FlowMatchEulerDiscreteScheduler,
|
|
)
|
|
|
|
|
|
def convert_vae(vae_ckpt_path):
|
|
old_state_dict = torch.load(vae_ckpt_path, weights_only=True)
|
|
new_state_dict = {}
|
|
|
|
middle_key_mapping = {
|
|
"encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma",
|
|
"encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias",
|
|
"encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight",
|
|
"encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma",
|
|
"encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias",
|
|
"encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight",
|
|
"encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma",
|
|
"encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias",
|
|
"encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight",
|
|
"encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma",
|
|
"encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias",
|
|
"encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight",
|
|
"decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma",
|
|
"decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias",
|
|
"decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight",
|
|
"decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma",
|
|
"decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias",
|
|
"decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight",
|
|
"decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma",
|
|
"decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias",
|
|
"decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight",
|
|
"decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma",
|
|
"decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias",
|
|
"decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight",
|
|
}
|
|
|
|
attention_mapping = {
|
|
"encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma",
|
|
"encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight",
|
|
"encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias",
|
|
"encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight",
|
|
"encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias",
|
|
"decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma",
|
|
"decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight",
|
|
"decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias",
|
|
"decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight",
|
|
"decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias",
|
|
}
|
|
|
|
head_mapping = {
|
|
"encoder.head.0.gamma": "encoder.norm_out.gamma",
|
|
"encoder.head.2.bias": "encoder.conv_out.bias",
|
|
"encoder.head.2.weight": "encoder.conv_out.weight",
|
|
"decoder.head.0.gamma": "decoder.norm_out.gamma",
|
|
"decoder.head.2.bias": "decoder.conv_out.bias",
|
|
"decoder.head.2.weight": "decoder.conv_out.weight",
|
|
}
|
|
|
|
quant_mapping = {
|
|
"conv1.weight": "quant_conv.weight",
|
|
"conv1.bias": "quant_conv.bias",
|
|
"conv2.weight": "post_quant_conv.weight",
|
|
"conv2.bias": "post_quant_conv.bias",
|
|
}
|
|
|
|
for key, value in old_state_dict.items():
|
|
if key in middle_key_mapping:
|
|
new_state_dict[middle_key_mapping[key]] = value
|
|
elif key in attention_mapping:
|
|
new_state_dict[attention_mapping[key]] = value
|
|
elif key in head_mapping:
|
|
new_state_dict[head_mapping[key]] = value
|
|
elif key in quant_mapping:
|
|
new_state_dict[quant_mapping[key]] = value
|
|
elif key == "encoder.conv1.weight":
|
|
new_state_dict["encoder.conv_in.weight"] = value
|
|
elif key == "encoder.conv1.bias":
|
|
new_state_dict["encoder.conv_in.bias"] = value
|
|
elif key == "decoder.conv1.weight":
|
|
new_state_dict["decoder.conv_in.weight"] = value
|
|
elif key == "decoder.conv1.bias":
|
|
new_state_dict["decoder.conv_in.bias"] = value
|
|
elif key.startswith("encoder.downsamples."):
|
|
new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.")
|
|
if ".residual.0.gamma" in new_key:
|
|
new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma")
|
|
elif ".residual.2.bias" in new_key:
|
|
new_key = new_key.replace(".residual.2.bias", ".conv1.bias")
|
|
elif ".residual.2.weight" in new_key:
|
|
new_key = new_key.replace(".residual.2.weight", ".conv1.weight")
|
|
elif ".residual.3.gamma" in new_key:
|
|
new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma")
|
|
elif ".residual.6.bias" in new_key:
|
|
new_key = new_key.replace(".residual.6.bias", ".conv2.bias")
|
|
elif ".residual.6.weight" in new_key:
|
|
new_key = new_key.replace(".residual.6.weight", ".conv2.weight")
|
|
elif ".shortcut.bias" in new_key:
|
|
new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias")
|
|
elif ".shortcut.weight" in new_key:
|
|
new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight")
|
|
new_state_dict[new_key] = value
|
|
elif key.startswith("decoder.upsamples."):
|
|
parts = key.split(".")
|
|
block_idx = int(parts[2])
|
|
|
|
if "residual" in key:
|
|
if block_idx in [0, 1, 2]:
|
|
new_block_idx = 0
|
|
resnet_idx = block_idx
|
|
elif block_idx in [4, 5, 6]:
|
|
new_block_idx = 1
|
|
resnet_idx = block_idx - 4
|
|
elif block_idx in [8, 9, 10]:
|
|
new_block_idx = 2
|
|
resnet_idx = block_idx - 8
|
|
elif block_idx in [12, 13, 14]:
|
|
new_block_idx = 3
|
|
resnet_idx = block_idx - 12
|
|
else:
|
|
new_state_dict[key] = value
|
|
continue
|
|
|
|
if ".residual.0.gamma" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma"
|
|
elif ".residual.2.bias" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias"
|
|
elif ".residual.2.weight" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight"
|
|
elif ".residual.3.gamma" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma"
|
|
elif ".residual.6.bias" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias"
|
|
elif ".residual.6.weight" in key:
|
|
new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight"
|
|
else:
|
|
new_key = key
|
|
new_state_dict[new_key] = value
|
|
|
|
elif ".shortcut." in key:
|
|
if block_idx == 4:
|
|
new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.")
|
|
new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1")
|
|
else:
|
|
new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.")
|
|
new_key = new_key.replace(".shortcut.", ".conv_shortcut.")
|
|
new_state_dict[new_key] = value
|
|
|
|
elif ".resample." in key or ".time_conv." in key:
|
|
if block_idx == 3:
|
|
new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0")
|
|
elif block_idx == 7:
|
|
new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0")
|
|
elif block_idx == 11:
|
|
new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0")
|
|
else:
|
|
new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.")
|
|
new_state_dict[new_key] = value
|
|
else:
|
|
new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.")
|
|
new_state_dict[new_key] = value
|
|
else:
|
|
new_state_dict[key] = value
|
|
|
|
with init_empty_weights():
|
|
vae = AutoencoderKLWan()
|
|
vae.load_state_dict(new_state_dict, strict=True, assign=True)
|
|
return vae
|
|
|
|
|
|
TRANSFORMER_CONFIG = {
|
|
"hidden_size": 4096,
|
|
"in_channels": 16,
|
|
"num_attention_heads": 32,
|
|
"num_layers": 40,
|
|
"out_channels": 16,
|
|
"patch_size": [1, 2, 2],
|
|
"rope_dim_list": [16, 56, 56],
|
|
"text_dim": 4096,
|
|
"rope_type": "rope",
|
|
"theta": 10000,
|
|
}
|
|
|
|
|
|
def convert_transformer(ckpt_path: str, model_type: str = "edit"):
|
|
checkpoint = torch.load(ckpt_path, weights_only=True)
|
|
if "model" in checkpoint:
|
|
original_state_dict = checkpoint["model"]
|
|
else:
|
|
original_state_dict = checkpoint
|
|
|
|
attn_suffixes = (
|
|
"img_attn_qkv.",
|
|
"img_attn_q_norm.",
|
|
"img_attn_k_norm.",
|
|
"img_attn_proj.",
|
|
"txt_attn_qkv.",
|
|
"txt_attn_q_norm.",
|
|
"txt_attn_k_norm.",
|
|
"txt_attn_proj.",
|
|
)
|
|
remapped = {}
|
|
for key, value in original_state_dict.items():
|
|
new_key = key
|
|
if key.startswith("double_blocks."):
|
|
for suffix in attn_suffixes:
|
|
if "." + suffix in key and ".attn." + suffix not in key:
|
|
new_key = key.replace("." + suffix, ".attn." + suffix)
|
|
break
|
|
remapped[new_key] = value
|
|
|
|
transformer_cls = (
|
|
JoyImageEditPlusTransformer3DModel if model_type == "edit_plus" else JoyImageEditTransformer3DModel
|
|
)
|
|
with init_empty_weights():
|
|
transformer = transformer_cls(**TRANSFORMER_CONFIG)
|
|
transformer.load_state_dict(remapped, strict=True, assign=True)
|
|
return transformer
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser(description="Convert JoyImage Edit / Edit Plus checkpoints to diffusers format")
|
|
parser.add_argument(
|
|
"--model_type",
|
|
type=str,
|
|
choices=["edit", "edit_plus"],
|
|
default="edit",
|
|
help="Model type: 'edit' for JoyImage-Edit, 'edit_plus' for JoyImage-Edit-Plus",
|
|
)
|
|
parser.add_argument(
|
|
"--transformer_ckpt_path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to original transformer checkpoint",
|
|
)
|
|
parser.add_argument(
|
|
"--vae_ckpt_path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to original VAE checkpoint",
|
|
)
|
|
parser.add_argument(
|
|
"--text_encoder_path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to Qwen3-VL text encoder (e.g. Qwen/Qwen3-VL-8B-Instruct)",
|
|
)
|
|
parser.add_argument("--save_pipeline", action="store_true")
|
|
parser.add_argument(
|
|
"--output_path",
|
|
type=str,
|
|
required=True,
|
|
help="Path where converted model should be saved",
|
|
)
|
|
parser.add_argument("--dtype", default="bf16", help="Torch dtype (fp32, fp16, bf16)")
|
|
parser.add_argument("--flow_shift", type=float, default=1.5)
|
|
return parser.parse_args()
|
|
|
|
|
|
DTYPE_MAPPING = {
|
|
"fp32": torch.float32,
|
|
"fp16": torch.float16,
|
|
"bf16": torch.bfloat16,
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
args = get_args()
|
|
transformer = None
|
|
vae = None
|
|
dtype = DTYPE_MAPPING[args.dtype]
|
|
|
|
if args.save_pipeline:
|
|
assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None
|
|
assert args.text_encoder_path is not None
|
|
|
|
if args.transformer_ckpt_path is not None:
|
|
transformer = convert_transformer(args.transformer_ckpt_path, model_type=args.model_type)
|
|
transformer = transformer.to(dtype=dtype)
|
|
if not args.save_pipeline:
|
|
transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
|
|
|
|
if args.vae_ckpt_path is not None:
|
|
vae = convert_vae(args.vae_ckpt_path)
|
|
vae = vae.to(dtype=dtype)
|
|
if not args.save_pipeline:
|
|
vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
|
|
|
|
if args.save_pipeline:
|
|
processor = AutoProcessor.from_pretrained(args.text_encoder_path)
|
|
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
|
|
args.text_encoder_path, torch_dtype=torch.bfloat16
|
|
).to("cuda")
|
|
tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path)
|
|
scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.flow_shift)
|
|
transformer = transformer.to("cuda")
|
|
vae = vae.to("cuda")
|
|
|
|
if args.model_type == "edit_plus":
|
|
pipe = JoyImageEditPlusPipeline(
|
|
processor=processor,
|
|
transformer=transformer,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
vae=vae,
|
|
scheduler=scheduler,
|
|
).to("cuda")
|
|
else:
|
|
pipe = JoyImageEditPipeline(
|
|
processor=processor,
|
|
transformer=transformer,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
vae=vae,
|
|
scheduler=scheduler,
|
|
).to("cuda")
|
|
|
|
pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
|
|
processor.save_pretrained(f"{args.output_path}/processor")
|