diff --git a/app.py b/app.py index 5be6b9e..04a526b 100644 --- a/app.py +++ b/app.py @@ -266,67 +266,102 @@ import torch.nn.functional as F -v21_path = hf_hub_download( - repo_id="Phr00t/Qwen-Image-Edit-Rapid-AIO", - filename="v21/Qwen-Rapid-AIO-NSFW-v21.safetensors", - repo_type="model" -) - -# 2. load the base architecture -# we use the default flowmatch scheduler first to ensure the pipe inits correctly, -# then we swap it to euler_a later print("loading base pipeline architecture...") pipe = QwenImageEditPlusPipeline.from_pretrained( "Qwen/Qwen-Image-Edit-2511", torch_dtype=torch.bfloat16 ).to("cuda") -# 3. switch scheduler to Euler Ancestral (Lightning requirement) -# we configure it with the base config to keep timestep spacing correct +# force euler ancestral scheduler pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) -# 4. load the massive 28GB v21 weights -print(f"loading v21 weights from {v21_path}...") +# 2. DOWNLOAD & LOAD RAW WEIGHTS +# ------------------------------------------------------------------------------ +print("accessing v21 checkpoint...") +v21_path = hf_hub_download( + repo_id="Phr00t/Qwen-Image-Edit-Rapid-AIO", + filename="v21/Qwen-Rapid-AIO-NSFW-v21.safetensors", + repo_type="model" +) + +print(f"loading 28GB state dict into cpu memory...") state_dict = load_file(v21_path) -# 5. The "Brutal" Injection -# Because this is an AIO file, keys might be prefixed with "model." or "transformer." -# or they might match the pipeline exactly. We try the root load first. -print("injecting AIO weights...") +# 3. DYNAMIC COMPONENT MAPPING (NO ASSUMPTIONS) +# ------------------------------------------------------------------------------ +print("sorting weights into components...") -# clean up keys if necessary (common in comfyui > diffusers conversions) -# this removes 'model.diffusion_model.' prefixes if they exist to match diffusers 'transformer.' -new_state_dict = {} +# containers for the sorted weights +transformer_weights = {} +vae_weights = {} +text_encoder_weights = {} + +# analyze the first key to determine the format +first_key = next(iter(state_dict.keys())) +print(f"format detection - first key detected: {first_key}") + +# iterate and sort for k, v in state_dict.items(): + # MAPPING: TRANSFORMER + # ComfyUI usually prefixes with 'model.diffusion_model.' if k.startswith("model.diffusion_model."): - new_key = k.replace("model.diffusion_model.", "transformer.") - new_state_dict[new_key] = v + new_key = k.replace("model.diffusion_model.", "") + transformer_weights[new_key] = v + # Or sometimes just 'transformer.' or 'model.' + elif k.startswith("transformer."): + new_key = k.replace("transformer.", "") + transformer_weights[new_key] = v + + # MAPPING: VAE + # ComfyUI prefix: 'first_stage_model.' elif k.startswith("first_stage_model."): - new_key = k.replace("first_stage_model.", "vae.") - new_state_dict[new_key] = v - elif k.startswith("conditioner.embedders.0."): - new_key = k.replace("conditioner.embedders.0.", "text_encoder.") - new_state_dict[new_key] = v - else: - new_state_dict[k] = v + new_key = k.replace("first_stage_model.", "") + vae_weights[new_key] = v + # Diffusers prefix: 'vae.' + elif k.startswith("vae."): + new_key = k.replace("vae.", "") + vae_weights[new_key] = v -# if no keys were renamed, just use the original -if len(new_state_dict) == len(state_dict): - final_dict = state_dict + # MAPPING: TEXT ENCODER + # ComfyUI prefix: 'conditioner.embedders.' or 'text_encoder.' + elif "text_encoder" in k or "conditioner" in k: + # this is tricky, we try to keep the suffix + if "conditioner.embedders.0." in k: + new_key = k.replace("conditioner.embedders.0.", "") + text_encoder_weights[new_key] = v + elif "text_encoder." in k: + new_key = k.replace("text_encoder.", "") + text_encoder_weights[new_key] = v + +# 4. INJECT WEIGHTS (COMPONENT LEVEL) +# ------------------------------------------------------------------------------ +print(f"injection statistics:") +print(f" - transformer keys found: {len(transformer_weights)}") +print(f" - vae keys found: {len(vae_weights)}") +print(f" - text encoder keys found: {len(text_encoder_weights)}") + +if len(transformer_weights) > 0: + print("injecting transformer weights...") + msg = pipe.transformer.load_state_dict(transformer_weights, strict=False) + print(f"transformer missing keys: {len(msg.missing_keys)}") else: - print("detected comfyui keys, remapped for diffusers.") - final_dict = new_state_dict + print("CRITICAL WARNING: no transformer weights found in file. check mapping logic.") -# attempt load -mismatched = pipe.load_state_dict(final_dict, strict=False) -print("weights loaded.") -print(f"missing keys (ignore if just config/aux): {len(mismatched.missing_keys)}") -print(f"unexpected keys (ignore if comfy artifacts): {len(mismatched.unexpected_keys)}") +if len(vae_weights) > 0: + print("injecting vae weights...") + pipe.vae.load_state_dict(vae_weights, strict=False) -# 6. cleanup +if len(text_encoder_weights) > 0: + print("injecting text encoder weights...") + # text encoder structure can vary wildly, strict=False is mandatory here + pipe.text_encoder.load_state_dict(text_encoder_weights, strict=False) + +# 5. CLEANUP & RUN +# ------------------------------------------------------------------------------ del state_dict -del new_state_dict -del final_dict +del transformer_weights +del vae_weights +del text_encoder_weights gc.collect() torch.cuda.empty_cache()