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")