Update app.py

app.py

@@ -254,113 +254,6 @@ pipe.fuse_lora()
 print("lightning lora fused.")
 
 
-# --- 2. manual surgery for lokr (snofs) ---
-print("attempting manual lokr injection for snofs...")
-
-try:
-    from safetensors.torch import load_file
-    from huggingface_hub import hf_hub_download
-    import torch.nn as nn
-    
-    # download
-    lora_path = hf_hub_download(repo_id="headlesssetton/kjfakjf", filename="Qwen_Snofs_1_2.safetensors")
-    state_dict = load_file(lora_path)
-    
-    lokr_scale = 1.0
-
-    def get_target_module(root, path_str):
-        """
-        recursively resolves a path like 'transformer_blocks_25_attn_add_v_proj'
-        against the actual model structure.
-        """
-        if not path_str:
-            return root
-            
-        # try to find the longest attribute/index match at the start of path_str
-        # we split by underscores, but we have to be careful about names that contain underscores
-        parts = path_str.split('_')
-        
-        # greedy search: try to consume as many parts as possible to form a valid attribute name
-        current_attr = ""
-        for i in range(len(parts), 0, -1):
-            candidate = "_".join(parts[:i])
-            
-            # check if candidate is an attribute
-            if hasattr(root, candidate):
-                next_root = getattr(root, candidate)
-                remaining = "_".join(parts[i:])
-                return get_target_module(next_root, remaining)
-            
-            # check if candidate is an index (for ModuleList/Sequential)
-            if candidate.isdigit() and isinstance(root, (nn.Sequential, nn.ModuleList)):
-                idx = int(candidate)
-                if idx < len(root):
-                    next_root = root[idx]
-                    remaining = "_".join(parts[i:])
-                    return get_target_module(next_root, remaining)
-        
-        # if we are here, we failed to match any attribute.
-        # heuristic: maybe the model uses 'proj' but lora says 'linear'? (unlikely here)
-        return None
-
-    updates = 0
-    with torch.no_grad():
-        prefixes = set()
-        for key in state_dict.keys():
-            if "lokr_w1" in key:
-                prefixes.add(key.replace(".lokr_w1", ""))
-        
-        for prefix in prefixes:
-            # clean prefix: remove 'lora_unet_'
-            search_path = prefix.replace("lora_unet_", "")
-            
-            # traverse
-            target_module = get_target_module(pipe.transformer, search_path)
-            
-            if target_module and hasattr(target_module, "weight"):
-                # load 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)
-                
-                # scale math
-                current_scale = lokr_scale
-                if alpha is not None:
-                    if isinstance(alpha, torch.Tensor):
-                        alpha = alpha.to(device, dtype=dtype)
-                    current_scale *= (alpha / w1.shape[0])
-
-                # kron
-                delta = torch.kron(w1, w2) * current_scale
-                
-                # device align
-                if delta.device != target_module.weight.device:
-                    delta = delta.to(target_module.weight.device)
-                
-                # inject
-                if target_module.weight.shape == delta.shape:
-                    target_module.weight.add_(delta)
-                    updates += 1
-                elif target_module.weight.shape == delta.T.shape:
-                    target_module.weight.add_(delta.T)
-                    updates += 1
-                else:
-                    print(f"shape mismatch {search_path}: {target_module.weight.shape} vs {delta.shape}")
-            else:
-                # keep this silent unless debugging, 120 success is usually enough for visual impact
-                pass
-
-    print(f"recursive solver successfully injected {updates} lokr layers.")
-
-except Exception as e:
-    import traceback
-    traceback.print_exc()
-    print(f"lokr injection failed: {e}")
-
-# --- end of surgery ---
-
-# --- end of surgery ---
-
 # # Apply the same optimizations from the first version
 # pipe.transformer.__class__ = QwenImageTransformer2DModel
 # pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())