ci/qwen-image
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update app.py

Browse Source
This commit is contained in:
sigma
2026-01-19 01:59:00 +00:00
committed by system
parent 6c17d7721a
commit 3bfb445e7b
1 changed files with 77 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
app.py
View File

@@ -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...") print("loading base pipeline architecture...")
pipe = QwenImageEditPlusPipeline.from_pretrained( pipe = QwenImageEditPlusPipeline.from_pretrained(
"Qwen/Qwen-Image-Edit-2511", "Qwen/Qwen-Image-Edit-2511",
torch_dtype=torch.bfloat16 torch_dtype=torch.bfloat16
).to("cuda") ).to("cuda")
# 3. switch scheduler to Euler Ancestral (Lightning requirement) # force euler ancestral scheduler
# we configure it with the base config to keep timestep spacing correct
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
# 4. load the massive 28GB v21 weights # 2. DOWNLOAD & LOAD RAW WEIGHTS
print(f"loading v21 weights from {v21_path}...") # ------------------------------------------------------------------------------
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) state_dict = load_file(v21_path)
# 5. The "Brutal" Injection # 3. DYNAMIC COMPONENT MAPPING (NO ASSUMPTIONS)
# 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("sorting weights into components...")
print("injecting AIO weights...")
# clean up keys if necessary (common in comfyui > diffusers conversions) # containers for the sorted weights
# this removes 'model.diffusion_model.' prefixes if they exist to match diffusers 'transformer.' transformer_weights = {}
new_state_dict = {} 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(): for k, v in state_dict.items():
# MAPPING: TRANSFORMER
# ComfyUI usually prefixes with 'model.diffusion_model.'
if k.startswith("model.diffusion_model."): if k.startswith("model.diffusion_model."):
new_key = k.replace("model.diffusion_model.", "transformer.") new_key = k.replace("model.diffusion_model.", "")
new_state_dict[new_key] = v 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."): elif k.startswith("first_stage_model."):
new_key = k.replace("first_stage_model.", "vae.") new_key = k.replace("first_stage_model.", "")
new_state_dict[new_key] = v vae_weights[new_key] = v
elif k.startswith("conditioner.embedders.0."): # Diffusers prefix: 'vae.'
new_key = k.replace("conditioner.embedders.0.", "text_encoder.") elif k.startswith("vae."):
new_state_dict[new_key] = v new_key = k.replace("vae.", "")
vae_weights[new_key] = v
# 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: else:
new_state_dict[k] = v print("CRITICAL WARNING: no transformer weights found in file. check mapping logic.")
# if no keys were renamed, just use the original if len(vae_weights) > 0:
if len(new_state_dict) == len(state_dict): print("injecting vae weights...")
final_dict = state_dict pipe.vae.load_state_dict(vae_weights, strict=False)
else:
print("detected comfyui keys, remapped for diffusers.")
final_dict = new_state_dict
# attempt load if len(text_encoder_weights) > 0:
mismatched = pipe.load_state_dict(final_dict, strict=False) print("injecting text encoder weights...")
print("weights loaded.") # text encoder structure can vary wildly, strict=False is mandatory here
print(f"missing keys (ignore if just config/aux): {len(mismatched.missing_keys)}") pipe.text_encoder.load_state_dict(text_encoder_weights, strict=False)
print(f"unexpected keys (ignore if comfy artifacts): {len(mismatched.unexpected_keys)}")
# 6. cleanup # 5. CLEANUP & RUN
# ------------------------------------------------------------------------------
del state_dict del state_dict
del new_state_dict del transformer_weights
del final_dict del vae_weights
del text_encoder_weights
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()