Update app.py

app.py

@@ -39,37 +39,51 @@ pipe.to("cuda")
 from huggingface_hub import hf_hub_download
 import torch
 
+from huggingface_hub import hf_hub_download
+import torch
+
 # Load FusionX enhancement LoRAs
 lora_adapters = []
 lora_weights = []
 
-# Helper to print model keys
-def print_model_keys(model, show_values=False):
-    print("\n🔑 Model State Dict Keys:")
-    for k, v in model.state_dict().items():
-        if show_values:
+# Print all named parameters (safely) from any pipeline
+def print_named_params(module, module_name=""):
+    print(f"\n🔍 Parameters in {module_name or 'pipeline'}:")
+    for name, param in module.named_parameters():
+        print(f"{name}: {param.shape}")
+
+# Try printing known submodules in the pipeline
+def print_all_pipeline_keys(pipe):
+    print("🧠 Exploring pipeline structure:")
+    for attr in dir(pipe):
+        if not attr.startswith("_"):
+            try:
+                obj = getattr(pipe, attr)
+                if isinstance(obj, torch.nn.Module):
+                    print_named_params(obj, attr)
+            except Exception as e:
+                print(f"⚠️ Could not inspect {attr}: {e}")
+
+# Print LoRA file contents
+def print_lora_checkpoint_keys(path):
+    try:
+        lora_checkpoint = torch.load(path, map_location="cpu")
+        print("\n📦 LoRA Checkpoint Keys:")
+        for k, v in lora_checkpoint.items():
             print(f"{k}: {v.shape}")
-        else:
-            print(k)
+    except Exception as e:
+        print(f"❌ Failed to load LoRA for inspection: {e}")
 
-# Print keys in UNet, Text Encoder, VAE
-print_model_keys(pipe.unet, show_values=True)
-print_model_keys(pipe.text_encoder, show_values=True)
-print_model_keys(pipe.vae, show_values=True)
+# Step 1: Explore the pipeline model structure
+print_all_pipeline_keys(pipe)
 
-# Download LoRA file
+# Step 2: Download LoRA file
 lora_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
 
-# Print keys from LoRA checkpoint
-print("\n📦 LoRA Checkpoint Keys:")
-try:
-    lora_checkpoint = torch.load(lora_path, map_location="cpu")
-    for k, v in lora_checkpoint.items():
-        print(f"{k}: {v.shape}")
-except Exception as e:
-    print(f"❌ Failed to load LoRA file for key inspection: {e}")
+# Step 3: Print LoRA checkpoint keys
+print_lora_checkpoint_keys(lora_path)
 
-# Load and fuse LoRA
+# Step 4: Load and apply the LoRA
 try:
     pipe.load_lora_weights(lora_path, adapter_name="main")
     pipe.set_adapters(["main"], adapter_weights=[1.0])