53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
import os
|
|
|
|
# Apply qwen_image21 detection patch
|
|
f = "/opt/ComfyUI/comfy/model_detection.py"
|
|
with open(f) as fh:
|
|
code = fh.read()
|
|
if "qwen_image21" not in code:
|
|
qwen21_check = (
|
|
" if '{}txt_in.text_norm.weight'.format(key_prefix) in state_dict_keys: # Qwen Image 2.1\n"
|
|
" dit_config = {}\n"
|
|
' dit_config["image_model"] = "qwen_image21"\n'
|
|
" dit_config[\"in_channels\"] = state_dict['{}img_in.weight'.format(key_prefix)].shape[0]\n"
|
|
" dit_config[\"num_layers\"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')\n"
|
|
" return dit_config\n"
|
|
)
|
|
anchor = " if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys: # Qwen Image"
|
|
code = code.replace(anchor, qwen21_check + "\n" + anchor)
|
|
with open(f, "w") as fh:
|
|
fh.write(code)
|
|
print("Applied qwen_image21 detection patch")
|
|
|
|
# Add qwen_image21 to GGUF node arch list
|
|
loader = "/opt/ComfyUI/custom_nodes/ComfyUI-GGUF/loader.py"
|
|
if os.path.exists(loader):
|
|
with open(loader) as fh:
|
|
lcode = fh.read()
|
|
if '"qwen_image21"' not in lcode:
|
|
old = '"lumina2", "qwen_image"}'
|
|
new = '"lumina2", "qwen_image", "qwen_image21"}'
|
|
if old in lcode:
|
|
lcode = lcode.replace(old, new)
|
|
with open(loader, "w") as fh:
|
|
fh.write(lcode)
|
|
print("Applied qwen_image21 arch list patch")
|
|
|
|
# Patch VAE decode to expand 4D latents to 5D for video VAEs (latent_dim=3)
|
|
sd = "/opt/ComfyUI/comfy/sd.py"
|
|
with open(sd) as fh:
|
|
sdc = fh.read()
|
|
if "self.latent_dim == 3 and samples_in.ndim == 4" not in sdc:
|
|
anchor = " if self.latent_dim == 2 and samples_in.ndim == 5:\n samples_in = samples_in[:, :, 0]"
|
|
replacement = (
|
|
" if self.latent_dim == 2 and samples_in.ndim == 5:\n"
|
|
" samples_in = samples_in[:, :, 0]\n"
|
|
" elif self.latent_dim == 3 and samples_in.ndim == 4:\n"
|
|
" samples_in = samples_in.unsqueeze(2)"
|
|
)
|
|
if anchor in sdc:
|
|
sdc = sdc.replace(anchor, replacement)
|
|
with open(sd, "w") as fh:
|
|
fh.write(sdc)
|
|
print("Applied VAE 4D-to-5D latent expansion patch")
|