Upload app.py with huggingface_hub

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 import random
 import numpy as np
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from huggingface_hub import hf_hub_download
 
 dtype = torch.bfloat16
 device = "cuda"
@@ -21,15 +22,56 @@ good_vae = AutoencoderKL.from_pretrained(
 # Load base FLUX.1-dev with tiny VAE
 BASE_MODEL = "black-forest-labs/FLUX.1-dev"
 CMO_LORA = "Bruece/FLUX.1-dev-CMO"
+LORA_ALPHA = 128
+LORA_R = 64
+LORA_SCALE = LORA_ALPHA / LORA_R  # = 2.0
 
 pipe = DiffusionPipeline.from_pretrained(
     BASE_MODEL, torch_dtype=dtype, vae=taef1
 ).to(device)
 
-# Load the CMO LoRA adapter using diffusers' built-in loader (ZeroGPU-safe)
-pipe.load_lora_weights(CMO_LORA)
-pipe.fuse_lora()
-
+# Manually load and merge the CMO LoRA adapter weights.
+# Use safetensors.safe_open (numpy backend, no torch) to avoid ZeroGPU's
+# torch patching which fails at module scope (no CUDA available yet).
+from safetensors import safe_open
+
+_lora_path = hf_hub_download(CMO_LORA, "adapter_model.safetensors", repo_type="model")
+
+# Load all LoRA weights as numpy arrays first, then merge into the transformer
+_lora_a_pairs = {}  # module_path -> A weight (numpy)
+_lora_b_pairs = {}  # module_path -> B weight (numpy)
+
+with safe_open(_lora_path, framework="numpy", device="cpu") as _f:
+    for _key in _f.keys():
+        if not _key.startswith("base_model.model."):
+            continue
+        _rest = _key[len("base_model.model."):]
+        if _rest.endswith(".lora_A.weight"):
+            _module_path = _rest[: -len(".lora_A.weight")]
+            _lora_a_pairs[_module_path] = _f.get_tensor(_key)
+        elif _rest.endswith(".lora_B.weight"):
+            _module_path = _rest[: -len(".lora_B.weight")]
+            _lora_b_pairs[_module_path] = _f.get_tensor(_key)
+
+# Merge LoRA weights into the transformer: w_new = w_orig + scale * (B @ A)
+_merge_count = 0
+for _module_path, _a_np in _lora_a_pairs.items():
+    if _module_path not in _lora_b_pairs:
+        continue
+    _b_np = _lora_b_pairs[_module_path]
+    # Navigate to the module in the transformer
+    _module = pipe.transformer
+    for _part in _module_path.split("."):
+        _module = getattr(_module, _part)
+    # Merge: w_orig + scale * (B @ A)
+    _a_tensor = torch.from_numpy(_a_np)
+    _b_tensor = torch.from_numpy(_b_np)
+    _delta = (_b_tensor.float() @ _a_tensor.float()) * LORA_SCALE
+    _module.weight.data.add_(_delta.to(_module.weight.data.dtype))
+    _merge_count += 1
+
+print(f"CMO LoRA: merged {_merge_count} adapter pairs into FLUX.1-dev transformer")
+del _lora_a_pairs, _lora_b_pairs
 torch.cuda.empty_cache()
 
 MAX_SEED = np.iinfo(np.int32).max