Update app.py

app.py

@@ -12,6 +12,7 @@ Lyra VAE Versions:
 """
 
 import os
+import json
 import torch
 import gradio as gr
 import numpy as np
@@ -901,17 +902,20 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
     print(f"🎵 Loading Lyra VAE v1 from {repo_id}...")
     
     try:
-        checkpoint_path = hf_hub_download(
-            repo_id=repo_id,
-            filename="best_model.pt",
-            repo_type="model"
-        )
-        
-        checkpoint = torch.load(checkpoint_path, map_location="cpu")
-        
-        if 'config' in checkpoint:
-            config_dict = checkpoint['config']
-        else:
+        # Try to download config.json first
+        try:
+            print("  📥 Downloading config.json...")
+            config_path = hf_hub_download(
+                repo_id=repo_id,
+                filename="config.json",
+                repo_type="model"
+            )
+            with open(config_path, 'r') as f:
+                config_dict = json.load(f)
+            print(f"  ✓ Config loaded: {config_dict.get('fusion_strategy', 'unknown')} fusion")
+        except Exception:
+            # Fallback to defaults if no config.json
+            print("  ⚠️ No config.json found, using defaults")
             config_dict = {
                 'modality_dims': {"clip": 768, "t5": 768},
                 'latent_dim': 768,
@@ -925,6 +929,24 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
                 'fusion_dropout': 0.1
             }
         
+        # Download model weights
+        print("  📥 Downloading model weights...")
+        try:
+            checkpoint_path = hf_hub_download(
+                repo_id=repo_id,
+                filename="model.pt",
+                repo_type="model"
+            )
+        except Exception:
+            # Fallback to best_model.pt
+            checkpoint_path = hf_hub_download(
+                repo_id=repo_id,
+                filename="best_model.pt",
+                repo_type="model"
+            )
+        
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        
         vae_config = LyraV1Config(
             modality_dims=config_dict.get('modality_dims', {"clip": 768, "t5": 768}),
             latent_dim=config_dict.get('latent_dim', 768),
@@ -948,11 +970,18 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
         lyra_model.to(device)
         lyra_model.eval()
         
-        print(f"✅ Lyra VAE v1 (SD1.5) loaded")
+        print(f"✅ Lyra VAE v1 loaded")
+        print(f"   Fusion: {config_dict.get('fusion_strategy')}")
+        print(f"   Latent dim: {config_dict.get('latent_dim')}")
+        if 'global_step' in checkpoint:
+            print(f"   Step: {checkpoint['global_step']:,}")
+        
         return lyra_model
         
     except Exception as e:
         print(f"❌ Failed to load Lyra VAE v1: {e}")
+        import traceback
+        traceback.print_exc()
         return None
 
 
@@ -968,46 +997,52 @@ def load_lyra_vae_xl(
     print(f"🎵 Loading Lyra VAE v2 from {repo_id}...")
     
     try:
+        # Download config.json first to get model architecture
+        print("  📥 Downloading config.json...")
+        config_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="config.json",
+            repo_type="model"
+        )
+        
+        with open(config_path, 'r') as f:
+            config_dict = json.load(f)
+        
+        print(f"  ✓ Config loaded: {config_dict.get('fusion_strategy', 'unknown')} fusion")
+        
+        # Download model weights
+        print("  📥 Downloading model.pt...")
         checkpoint_path = hf_hub_download(
             repo_id=repo_id,
-            filename="best_model.pt",
+            filename="model.pt",
             repo_type="model"
         )
         
         checkpoint = torch.load(checkpoint_path, map_location="cpu")
         
-        if 'config' in checkpoint:
-            config_dict = checkpoint['config']
-        else:
-            # XL v2 defaults - larger dimensions for SDXL
-            config_dict = {
-                'modality_dims': {"clip": 768, "t5": 2048},  # T5-XL
-                'latent_dim': 2048,
-                'seq_len': 77,
-                'encoder_layers': 4,
-                'decoder_layers': 4,
-                'hidden_dim': 2048,
-                'dropout': 0.1,
-                'fusion_strategy': 'adaptive_cantor',
-                'fusion_heads': 16,
-                'fusion_dropout': 0.1
-            }
-        
+        # Build config from repo's config.json
         vae_config = LyraV2Config(
-            modality_dims=config_dict.get('modality_dims', {"clip": 768, "t5": 2048}),
+            modality_dims=config_dict.get('modality_dims', {"clip_l": 768, "clip_g": 1280, "t5_xl_l": 2048, "t5_xl_g": 2048}),
+            modality_seq_lens=config_dict.get('modality_seq_lens', {"clip_l": 77, "clip_g": 77, "t5_xl_l": 512, "t5_xl_g": 512}),
+            binding_config=config_dict.get('binding_config'),
             latent_dim=config_dict.get('latent_dim', 2048),
             seq_len=config_dict.get('seq_len', 77),
-            encoder_layers=config_dict.get('encoder_layers', 4),
-            decoder_layers=config_dict.get('decoder_layers', 4),
+            encoder_layers=config_dict.get('encoder_layers', 3),
+            decoder_layers=config_dict.get('decoder_layers', 3),
             hidden_dim=config_dict.get('hidden_dim', 2048),
             dropout=config_dict.get('dropout', 0.1),
             fusion_strategy=config_dict.get('fusion_strategy', 'adaptive_cantor'),
-            fusion_heads=config_dict.get('fusion_heads', 16),
-            fusion_dropout=config_dict.get('fusion_dropout', 0.1)
+            fusion_heads=config_dict.get('fusion_heads', 8),
+            fusion_dropout=config_dict.get('fusion_dropout', 0.1),
+            cantor_depth=config_dict.get('cantor_depth', 8),
+            cantor_local_window=config_dict.get('cantor_local_window', 3),
+            alpha_init=config_dict.get('alpha_init', 1.0),
+            beta_init=config_dict.get('beta_init', 0.3),
         )
         
         lyra_model = LyraV2(vae_config)
         
+        # Load weights from checkpoint
         if 'model_state_dict' in checkpoint:
             lyra_model.load_state_dict(checkpoint['model_state_dict'])
         else:
@@ -1016,14 +1051,21 @@ def load_lyra_vae_xl(
         lyra_model.to(device)
         lyra_model.eval()
         
-        print(f"✅ Lyra VAE v2 (SDXL) loaded")
+        print(f"✅ Lyra VAE v2 loaded")
+        print(f"   Fusion: {config_dict.get('fusion_strategy')}")
+        print(f"   Latent dim: {config_dict.get('latent_dim')}")
+        print(f"   Hidden dim: {config_dict.get('hidden_dim')}")
         if 'global_step' in checkpoint:
             print(f"   Step: {checkpoint['global_step']:,}")
+        if 'best_loss' in checkpoint:
+            print(f"   Loss: {checkpoint['best_loss']:.4f}")
         
         return lyra_model
         
     except Exception as e:
         print(f"❌ Failed to load Lyra VAE v2: {e}")
+        import traceback
+        traceback.print_exc()
         return None