diff --git a/app.py b/app.py index b895ec5..ff387d1 100644 --- a/app.py +++ b/app.py @@ -234,20 +234,89 @@ scheduler_config = { scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config) # Load the model pipeline -pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511", - scheduler=scheduler, - torch_dtype=dtype).to(device) +from safetensors.torch import load_file +from huggingface_hub import hf_hub_download +import torch.nn.functional as F + +# --- 1. setup pipeline with lightning (this works fine) --- +pipe = QwenImageEditPlusPipeline.from_pretrained( + "Qwen/Qwen-Image-Edit-2511", + scheduler=scheduler, + torch_dtype=dtype +).to(device) + +print("loading lightning lora...") pipe.load_lora_weights( - "lightx2v/Qwen-Image-Edit-2511-Lightning", - weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors" + "lightx2v/Qwen-Image-Edit-2511-Lightning", + weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors" ) pipe.fuse_lora() +print("lightning lora fused.") -pipe.load_lora_weights( - "headlesssetton/kjfakjf", - weight_name="Qwen_Snofs_1_2.safetensors" -) -pipe.fuse_lora(lora_scale=1.0) +# --- 2. manual surgery for lokr (snofs) --- +print("attempting manual lokr injection for snofs...") + +try: + # download the file directly + lora_path = hf_hub_download(repo_id="headlesssetton/kjfakjf", filename="Qwen_Snofs_1_2.safetensors") + state_dict = load_file(lora_path) + + # lokr injection parameters + lokr_scale = 1.0 # adjust strength here + + # iterate and merge + updates = 0 + with torch.no_grad(): + # group keys by layer prefix + prefixes = set() + for key in state_dict.keys(): + if "lokr_w1" in key: + prefixes.add(key.replace(".lokr_w1", "")) + + for prefix in prefixes: + # extract weights + w1 = state_dict[f"{prefix}.lokr_w1"].to(device, dtype=dtype) + w2 = state_dict[f"{prefix}.lokr_w2"].to(device, dtype=dtype) + alpha = state_dict.get(f"{prefix}.alpha", None) + + # calculate scaling + # lokr usually uses alpha / sqrt(rank) or similar, but often just alpha is enough + # if alpha is present, scale = alpha / w1.shape[0] (or similar convention) + # here we will assume simple multiplication or alpha scaling if provided + scale = lokr_scale + if alpha is not None: + scale *= (alpha / w1.shape[0]) # standard lora scaling convention, might vary for lokr + + # compute delta: kronecker product + # w1: (a, b), w2: (c, d) -> result: (a*c, b*d) + # torch.kron is (a*c, b*d) + delta = torch.kron(w1, w2) * scale + + # find target layer in model + # prefix example: "transformer_blocks.0.attn.add_k_proj" + # pipe.transformer matches this structure directly + path_parts = prefix.split('.') + target = pipe.transformer + try: + for part in path_parts: + target = getattr(target, part) + + # check shapes + if target.weight.shape == delta.shape: + target.weight.add_(delta) # in-place merge + updates += 1 + else: + print(f"shape mismatch for {prefix}: model {target.weight.shape} vs lora {delta.shape}") + except AttributeError: + print(f"layer not found: {prefix}") + + print(f"successfully injected {updates} lokr layers manually.") + +except Exception as e: + print(f"lokr injection failed: {e}") + print("running with lightning lora only.") + +# --- end of surgery --- # # Apply the same optimizations from the first version # pipe.transformer.__class__ = QwenImageTransformer2DModel