Update app.py

app.py

@@ -5,6 +5,10 @@ License: MIT
 
 SD1.5 and SDXL-based flow matching with geometric crystalline architectures.
 Supports Illustrious XL, standard SDXL, and SD1.5 variants.
+
+Lyra VAE Versions:
+- v1: SD1.5 (768 dim CLIP + T5-base) - geofractal.model.vae.vae_lyra
+- v2: SDXL/Illustrious (768 CLIP-L + 2048 T5-XL) - geofractal.model.vae.vae_lyra_v2
 """
 
 import os
@@ -32,9 +36,21 @@ from transformers import (
 )
 from huggingface_hub import hf_hub_download
 
-# Import Lyra VAE from geofractal
+# Import Lyra VAE v1 (SD1.5) from geofractal
+try:
+    from geofractal.model.vae.vae_lyra import MultiModalVAE as LyraV1, MultiModalVAEConfig as LyraV1Config
+    LYRA_V1_AVAILABLE = True
+except ImportError:
+    print("⚠️ Lyra VAE v1 not available")
+    LYRA_V1_AVAILABLE = False
 
-from geofractal.model.vae.vae_lyra_v2 import MultiModalVAE, MultiModalVAEConfig
+# Import Lyra VAE v2 (SDXL/Illustrious) from geofractal
+try:
+    from geofractal.model.vae.vae_lyra_v2 import MultiModalVAE as LyraV2, MultiModalVAEConfig as LyraV2Config
+    LYRA_V2_AVAILABLE = True
+except ImportError:
+    print("⚠️ Lyra VAE v2 not available")
+    LYRA_V2_AVAILABLE = False
 
 
 # ============================================================================
@@ -877,12 +893,12 @@ def load_sdxl_base(device: str = "cuda"):
 
 
 def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
-    """Load Lyra VAE (SD1.5 version) from HuggingFace."""
-    if not LYRA_AVAILABLE:
-        print("⚠️ Lyra VAE not available")
+    """Load Lyra VAE v1 (SD1.5 version) from HuggingFace."""
+    if not LYRA_V1_AVAILABLE:
+        print("⚠️ Lyra VAE v1 not available")
         return None
     
-    print(f"🎵 Loading Lyra VAE from {repo_id}...")
+    print(f"🎵 Loading Lyra VAE v1 from {repo_id}...")
     
     try:
         checkpoint_path = hf_hub_download(
@@ -909,7 +925,7 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
                 'fusion_dropout': 0.1
             }
         
-        vae_config = MultiModalVAEConfig(
+        vae_config = LyraV1Config(
             modality_dims=config_dict.get('modality_dims', {"clip": 768, "t5": 768}),
             latent_dim=config_dict.get('latent_dim', 768),
             seq_len=config_dict.get('seq_len', 77),
@@ -922,7 +938,7 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
             fusion_dropout=config_dict.get('fusion_dropout', 0.1)
         )
         
-        lyra_model = MultiModalVAE(vae_config)
+        lyra_model = LyraV1(vae_config)
         
         if 'model_state_dict' in checkpoint:
             lyra_model.load_state_dict(checkpoint['model_state_dict'])
@@ -932,11 +948,11 @@ def load_lyra_vae(repo_id: str = "AbstractPhil/vae-lyra", device: str = "cuda"):
         lyra_model.to(device)
         lyra_model.eval()
         
-        print(f"✅ Lyra VAE (SD1.5) loaded")
+        print(f"✅ Lyra VAE v1 (SD1.5) loaded")
         return lyra_model
         
     except Exception as e:
-        print(f"❌ Failed to load Lyra VAE: {e}")
+        print(f"❌ Failed to load Lyra VAE v1: {e}")
         return None
 
 
@@ -944,12 +960,12 @@ def load_lyra_vae_xl(
     repo_id: str = "AbstractPhil/vae-lyra-xl-adaptive-cantor-illustrious",
     device: str = "cuda"
 ):
-    """Load Lyra VAE XL version for SDXL/Illustrious."""
-    if not LYRA_AVAILABLE:
-        print("⚠️ Lyra VAE not available")
+    """Load Lyra VAE v2 (SDXL/Illustrious version) from HuggingFace."""
+    if not LYRA_V2_AVAILABLE:
+        print("⚠️ Lyra VAE v2 not available")
         return None
     
-    print(f"🎵 Loading Lyra VAE XL from {repo_id}...")
+    print(f"🎵 Loading Lyra VAE v2 from {repo_id}...")
     
     try:
         checkpoint_path = hf_hub_download(
@@ -963,7 +979,7 @@ def load_lyra_vae_xl(
         if 'config' in checkpoint:
             config_dict = checkpoint['config']
         else:
-            # XL defaults - note larger dimensions
+            # XL v2 defaults - larger dimensions for SDXL
             config_dict = {
                 'modality_dims': {"clip": 768, "t5": 2048},  # T5-XL
                 'latent_dim': 2048,
@@ -977,7 +993,7 @@ def load_lyra_vae_xl(
                 'fusion_dropout': 0.1
             }
         
-        vae_config = MultiModalVAEConfig(
+        vae_config = LyraV2Config(
             modality_dims=config_dict.get('modality_dims', {"clip": 768, "t5": 2048}),
             latent_dim=config_dict.get('latent_dim', 2048),
             seq_len=config_dict.get('seq_len', 77),
@@ -990,7 +1006,7 @@ def load_lyra_vae_xl(
             fusion_dropout=config_dict.get('fusion_dropout', 0.1)
         )
         
-        lyra_model = MultiModalVAE(vae_config)
+        lyra_model = LyraV2(vae_config)
         
         if 'model_state_dict' in checkpoint:
             lyra_model.load_state_dict(checkpoint['model_state_dict'])
@@ -1000,14 +1016,14 @@ def load_lyra_vae_xl(
         lyra_model.to(device)
         lyra_model.eval()
         
-        print(f"✅ Lyra VAE XL loaded")
+        print(f"✅ Lyra VAE v2 (SDXL) loaded")
         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 XL: {e}")
+        print(f"❌ Failed to load Lyra VAE v2: {e}")
         return None
 
 
@@ -1284,12 +1300,12 @@ def create_demo():
         
         Generate images using SD1.5 and SDXL-based models with geometric deep learning:
         
-        | Model | Architecture | Best For |
-        |-------|-------------|----------|
-        | **Illustrious XL** | SDXL | Anime/illustration, high detail |
-        | **SDXL Base** | SDXL | Photorealistic, general purpose |
-        | **Flow-Lune** | SD1.5 | Fast flow matching (15-25 steps) |
-        | **SD1.5 Base** | SD1.5 | Baseline comparison |
+        | Model | Architecture | Lyra Version | Best For |
+        |-------|-------------|--------------|----------|
+        | **Illustrious XL** | SDXL | v2 (T5-XL) | Anime/illustration, high detail |
+        | **SDXL Base** | SDXL | v2 (T5-XL) | Photorealistic, general purpose |
+        | **Flow-Lune** | SD1.5 | v1 (T5-base) | Fast flow matching (15-25 steps) |
+        | **SD1.5 Base** | SD1.5 | v1 (T5-base) | Baseline comparison |
         
         Enable **Lyra VAE** for CLIP+T5 fusion comparison!
         """)
@@ -1417,12 +1433,13 @@ def create_demo():
                 - **Illustrious XL**: Use CLIP skip 2, booru-style tags
                 - **SDXL Base**: Natural language prompts work well
                 - **Flow-Lune**: Enable flow matching, shift ~2.5, fewer steps
-                - **Lyra**: Generates both standard and fused for comparison
+                - **Lyra v2**: SDXL models use T5-XL for richer semantics
+                - **Lyra v1**: SD1.5 models use T5-base
                 
                 ### Model Info
                 - SDXL models use **epsilon** prediction
                 - Lune uses **v_prediction** with flow matching
-                - Lyra fuses CLIP + T5 for richer semantics
+                - Lyra fuses CLIP + T5 via geometric Cantor attention
                 """)
         
         # Examples