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[1]\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")