AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity
Files
remove-lora-utils-tests
diffusers/scripts/convert_joyimage_edit_to_diffusers.py
tangyanf 408ce8db29 feat: add JoyImage edit plus (#14032) 
* 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>
2026-07-07 17:45:06 -07:00

356 lines
15 KiB
Python
Raw Permalink Blame History

"""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")
View Git Blame Copy Permalink