Update app.py

app.py

@@ -327,20 +327,29 @@ class LazyLyraModel:
         # Config already prefetched
         config_dict = self._config
         
-        # Find checkpoint
-        from huggingface_hub import list_repo_files
-        
-        repo_files = list_repo_files(self.repo_id, repo_type="model")
-        checkpoint_files = [f for f in repo_files if f.endswith('.safetensors') or f.endswith('.pt')]
-        
-        # Prefer weights/ folder
-        weights_files = [f for f in checkpoint_files if f.startswith('weights/')]
-        if weights_files:
-            checkpoint_file = sorted(weights_files)[-1]  # Latest
-        elif checkpoint_files:
-            checkpoint_file = checkpoint_files[0]
+        # Use provided checkpoint or find one
+        if self.checkpoint and self.checkpoint.strip():
+            checkpoint_file = self.checkpoint.strip()
+            # Add weights/ prefix if not present and file doesn't exist at root
+            if not checkpoint_file.startswith('weights/'):
+                checkpoint_file = f"weights/{checkpoint_file}"
+            print(f"[Lyra] Using specified checkpoint: {checkpoint_file}")
         else:
-            raise FileNotFoundError(f"No checkpoint found in {self.repo_id}")
+            # Find checkpoint automatically
+            from huggingface_hub import list_repo_files
+            
+            repo_files = list_repo_files(self.repo_id, repo_type="model")
+            checkpoint_files = [f for f in repo_files if f.endswith('.safetensors') or f.endswith('.pt')]
+            
+            # Prefer weights/ folder
+            weights_files = [f for f in checkpoint_files if f.startswith('weights/')]
+            if weights_files:
+                checkpoint_file = sorted(weights_files)[-1]  # Latest
+            elif checkpoint_files:
+                checkpoint_file = checkpoint_files[0]
+            else:
+                raise FileNotFoundError(f"No checkpoint found in {self.repo_id}")
+            print(f"[Lyra] Auto-selected checkpoint: {checkpoint_file}")
         
         checkpoint_path = hf_hub_download(
             repo_id=self.repo_id,
@@ -1295,12 +1304,16 @@ def load_lune_checkpoint(repo_id: str, filename: str, device: str = "cuda"):
 
 def load_illustrious_xl(
     repo_id: str = "AbstractPhil/vae-lyra-xl-adaptive-cantor-illustrious",
-    filename: str = "illustriousXL_v01.safetensors",
+    filename: str = "",
     device: str = "cuda"
 ) -> Tuple[UNet2DConditionModel, AutoencoderKL, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer, CLIPTokenizer]:
     """Load Illustrious XL from single safetensors file."""
     from diffusers import StableDiffusionXLPipeline
     
+    # Default checkpoint if none specified
+    if not filename or not filename.strip():
+        filename = "illustriousXL_v01.safetensors"
+    
     print(f"📥 Loading Illustrious XL: {repo_id}/{filename}")
     
     checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
@@ -1377,13 +1390,23 @@ def load_sdxl_base(device: str = "cuda"):
 # PIPELINE INITIALIZATION
 # ============================================================================
 
-def initialize_pipeline(model_choice: str, device: str = "cuda"):
+def initialize_pipeline(model_choice: str, device: str = "cuda", checkpoint: str = "", lyra_checkpoint: str = ""):
     """Initialize the complete pipeline based on model choice.
     
     Uses lazy loading for T5 and Lyra - they won't be downloaded until first use.
+    
+    Args:
+        model_choice: Model selection from dropdown
+        device: Target device
+        checkpoint: Optional custom checkpoint filename (e.g., "my_model.safetensors")
+        lyra_checkpoint: Optional custom Lyra VAE checkpoint filename
     """
     
     print(f"🚀 Initializing {model_choice} pipeline...")
+    if checkpoint:
+        print(f"   Custom model checkpoint: {checkpoint}")
+    if lyra_checkpoint:
+        print(f"   Custom Lyra checkpoint: {lyra_checkpoint}")
     
     is_sdxl = "Illustrious" in model_choice or "SDXL" in model_choice
     is_lune = "Lune" in model_choice
@@ -1391,7 +1414,10 @@ def initialize_pipeline(model_choice: str, device: str = "cuda"):
     if is_sdxl:
         # SDXL-based models
         if "Illustrious" in model_choice:
-            unet, vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2 = load_illustrious_xl(device=device)
+            unet, vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2 = load_illustrious_xl(
+                device=device,
+                filename=checkpoint
+            )
         else:
             unet, vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2 = load_sdxl_base(device=device)
         
@@ -1404,7 +1430,8 @@ def initialize_pipeline(model_choice: str, device: str = "cuda"):
         )
         lyra_loader = LazyLyraModel(
             repo_id=LYRA_ILLUSTRIOUS_REPO,
-            device=device
+            device=device,
+            checkpoint=lyra_checkpoint if lyra_checkpoint and lyra_checkpoint.strip() else None
         )
         
         # Default scheduler: Euler Ancestral
@@ -1448,13 +1475,18 @@ def initialize_pipeline(model_choice: str, device: str = "cuda"):
         )
         lyra_loader = LazyLyraModel(
             repo_id=LYRA_SD15_REPO,
-            device=device
+            device=device,
+            checkpoint=lyra_checkpoint if lyra_checkpoint and lyra_checkpoint.strip() else None
         )
         
         # Load UNet
         if is_lune:
             repo_id = "AbstractPhil/sd15-flow-lune"
-            filename = "sd15_flow_lune_e34_s34000.pt"
+            # Use custom checkpoint or default
+            if checkpoint and checkpoint.strip():
+                filename = checkpoint
+            else:
+                filename = "sd15_flow_lune_e34_s34000.pt"
             unet = load_lune_checkpoint(repo_id, filename, device)
         else:
             unet = UNet2DConditionModel.from_pretrained(
@@ -1491,15 +1523,31 @@ def initialize_pipeline(model_choice: str, device: str = "cuda"):
 
 CURRENT_PIPELINE = None
 CURRENT_MODEL = None
+CURRENT_CHECKPOINT = None
+CURRENT_LYRA_CHECKPOINT = None
 
 
-def get_pipeline(model_choice: str):
+def get_pipeline(model_choice: str, checkpoint: str = "", lyra_checkpoint: str = ""):
     """Get or create pipeline for selected model."""
-    global CURRENT_PIPELINE, CURRENT_MODEL
-    
-    if CURRENT_PIPELINE is None or CURRENT_MODEL != model_choice:
-        CURRENT_PIPELINE = initialize_pipeline(model_choice, device="cuda")
+    global CURRENT_PIPELINE, CURRENT_MODEL, CURRENT_CHECKPOINT, CURRENT_LYRA_CHECKPOINT
+    
+    # Normalize empty values
+    checkpoint = checkpoint.strip() if checkpoint else ""
+    lyra_checkpoint = lyra_checkpoint.strip() if lyra_checkpoint else ""
+    
+    # Reinitialize if model or any checkpoint changed
+    if (CURRENT_PIPELINE is None or 
+        CURRENT_MODEL != model_choice or 
+        CURRENT_CHECKPOINT != checkpoint or
+        CURRENT_LYRA_CHECKPOINT != lyra_checkpoint):
+        CURRENT_PIPELINE = initialize_pipeline(
+            model_choice, device="cuda", 
+            checkpoint=checkpoint, 
+            lyra_checkpoint=lyra_checkpoint
+        )
         CURRENT_MODEL = model_choice
+        CURRENT_CHECKPOINT = checkpoint
+        CURRENT_LYRA_CHECKPOINT = lyra_checkpoint
     
     return CURRENT_PIPELINE
 
@@ -1522,7 +1570,7 @@ def estimate_duration(num_steps: int, width: int, height: int, use_lyra: bool =
 
 
 @spaces.GPU(duration=lambda *args: estimate_duration(
-    args[6], args[8], args[9], args[12], 
+    args[8], args[10], args[11], args[14], 
     "SDXL" in args[3] or "Illustrious" in args[3]
 ))
 def generate_image(
@@ -1530,6 +1578,8 @@ def generate_image(
     t5_summary: str,
     negative_prompt: str,
     model_choice: str,
+    checkpoint: str,
+    lyra_checkpoint: str,
     scheduler_choice: str,
     clip_skip: int,
     num_steps: int,
@@ -1551,6 +1601,8 @@ def generate_image(
     Args:
         prompt: Tags/keywords (CLIP input)
         t5_summary: Natural language summary (T5 input, unless clip_include_summary)
+        checkpoint: Custom model checkpoint filename (empty for default)
+        lyra_checkpoint: Custom Lyra VAE checkpoint filename (empty for default)
         use_separator: Use ¶ separator between tags and summary
         clip_include_summary: If True, CLIP also sees the summary
     """
@@ -1567,7 +1619,7 @@ def generate_image(
         progress((step + 1) / total, desc=desc)
     
     try:
-        pipeline = get_pipeline(model_choice)
+        pipeline = get_pipeline(model_choice, checkpoint, lyra_checkpoint)
         
         # Update scheduler if needed (SDXL only)
         is_sdxl = "SDXL" in model_choice or "Illustrious" in model_choice
@@ -1710,6 +1762,21 @@ def create_demo():
                     value="Illustrious XL"
                 )
                 
+                with gr.Accordion("Advanced Options", open=False):
+                    checkpoint = gr.Textbox(
+                        label="Model Checkpoint (optional)",
+                        value="",
+                        placeholder="e.g., illustriousXL_v01.safetensors",
+                        info="Leave empty for default. Illustrious: .safetensors, Lune: .pt"
+                    )
+                    
+                    lyra_checkpoint = gr.Textbox(
+                        label="Lyra VAE Checkpoint (optional)",
+                        value="weights/lyra_illustrious_step_41000.safetensors",
+                        placeholder="e.g., lyra_e100_s50000.safetensors",
+                        info="Leave empty for latest. Loaded from weights/ folder in Lyra repo."
+                    )
+                
                 scheduler_choice = gr.Dropdown(
                     label="Scheduler (SDXL only)",
                     choices=SCHEDULER_CHOICES,
@@ -1913,7 +1980,8 @@ def create_demo():
         generate_btn.click(
             fn=generate_image,
             inputs=[
-                prompt, t5_summary, negative_prompt, model_choice, scheduler_choice, clip_skip,
+                prompt, t5_summary, negative_prompt, model_choice, checkpoint, lyra_checkpoint,
+                scheduler_choice, clip_skip,
                 num_steps, cfg_scale, width, height, shift,
                 use_flow_matching, use_lyra, lyra_strength, use_separator, clip_include_summary,
                 seed, randomize_seed