Add cancel button, validate cache files, detect cache hits, update peft dependency, fix race conditions, enhance feedback, improve interrupt handling.

app.py

@@ -361,7 +361,6 @@ def create_app():
             return (
                 '<div class="status-warning">⏳ Loading started...</div>',
                 "Starting download...",
-                gr.update(interactive=False)
             )
 
         def on_load_pipeline_complete(status_msg, progress_text):
@@ -385,14 +384,26 @@ def create_app():
                     gr.update(interactive=True)
                 )
 
+        # Cancel button for pipeline loading
+        cancel_load_btn = gr.Button("🛑 Cancel Loading", variant="secondary", size="sm", visible=False)
+
         load_btn.click(
             fn=on_load_pipeline_start,
             inputs=[],
-            outputs=[load_status, load_progress, load_btn],
+            outputs=[load_status, load_progress],
+        ).then(
+            fn=lambda: (gr.update(visible=True), gr.update(interactive=False)),
+            inputs=[],
+            outputs=[cancel_load_btn, load_btn],
         ).then(
             fn=load_pipeline,
             inputs=[checkpoint_url, vae_url, lora_urls, lora_strengths],
             outputs=[load_status, load_progress],
+            show_progress="full",
+        ).then(
+            fn=lambda: (gr.update(visible=False), gr.update(interactive=True)),
+            inputs=[],
+            outputs=[cancel_load_btn, load_btn],
         ).then(
             fn=on_load_pipeline_complete,
             inputs=[load_status, load_progress],
@@ -407,6 +418,17 @@ def create_app():
             outputs=[cached_checkpoints, cached_vaes, cached_loras],
         )
 
+        # Cancel button handler
+        cancel_load_btn.click(
+            fn=lambda: (cancel_download(), 
+                       '<div class="status-warning">⏳ Cancelling...</div>',
+                       "Cancelling download...",
+                       gr.update(visible=False),
+                       gr.update(interactive=True)),
+            inputs=[],
+            outputs=[load_status, load_progress, cancel_load_btn, load_btn],
+        )
+
         def on_cached_checkpoint_change(cached_path):
             """Update URL when a cached checkpoint is selected."""
             if cached_path and cached_path != "(None found)":

requirements.txt

@@ -7,6 +7,9 @@ transformers>=4.35.0
 safetensors>=0.4.0
 optimum>=1.0.0
 
+# LoRA support (required for diffusers >= 0.26)
+peft>=0.6.0
+
 # Image processing
 Pillow>=10.0.0
 

src/config.py

@@ -133,3 +133,65 @@ def get_cached_loras():
     for file in sorted(CACHE_DIR.glob("*_lora.safetensors")):
         models.append(str(file))
     return models
+
+
+def validate_cache_file(cache_path: Path, min_size_mb: float = 1.0) -> tuple[bool, str]:
+    """
+    Validate a cached model file exists and has valid content.
+
+    Args:
+        cache_path: Path to the cached .safetensors file
+        min_size_mb: Minimum acceptable file size in MB (default: 1MB)
+
+    Returns:
+        Tuple of (is_valid, message)
+        - is_valid: True if file passes all checks
+        - message: Description of validation result
+    """
+    try:
+        if not cache_path.exists():
+            return False, f"File does not exist: {cache_path.name}"
+
+        if not cache_path.is_file():
+            return False, f"Not a regular file: {cache_path.name}"
+
+        file_size = cache_path.stat().st_size
+        size_mb = file_size / (1024 * 1024)
+
+        if size_mb < min_size_mb:
+            return False, f"File too small ({size_mb:.2f} MB < {min_size_mb} MB): {cache_path.name}"
+
+        # Check if it's a valid safetensors file by reading the header
+        if not cache_path.suffix == ".safetensors":
+            return True, f"Valid non-safetensors file: {cache_path.name}"
+
+        try:
+            with open(cache_path, "rb") as f:
+                # Read first 8 bytes (header size)
+                header_size_bytes = f.read(8)
+                if len(header_size_bytes) < 8:
+                    return False, f"File too small for safetensors header: {cache_path.name}"
+
+                import struct
+                header_size = struct.unpack("<Q", header_size_bytes)[0]
+
+                if header_size == 0:
+                    return False, f"Invalid safetensors header (size=0): {cache_path.name}"
+
+                # Read and parse header JSON
+                header = f.read(header_size)
+                if len(header) < header_size:
+                    return False, f"Incomplete safetensors header: {cache_path.name}"
+
+                import json
+                json.loads(header.decode("utf-8"))
+
+        except struct.error as e:
+            return False, f"Invalid safetensors format: {str(e)}"
+        except json.JSONDecodeError as e:
+            return False, f"Invalid safetensors header JSON: {str(e)}"
+
+        return True, f"Valid cached file ({size_mb:.1f} MB): {cache_path.name}"
+
+    except OSError as e:
+        return False, f"File access error: {str(e)}"

src/downloader.py

@@ -219,6 +219,9 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         if local_path.exists():
             import shutil
             output_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            print(f"  📁 Copying from cache: {local_path.name} → {output_path.name}")
+            
             # Copy the file to cache location
             shutil.copy2(str(local_path), str(output_path))
 
@@ -229,6 +232,8 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         else:
             raise FileNotFoundError(f"Local file not found: {local_path}")
 
+    print(f"  📥 Downloading to cache: {output_path.name}")
+    
     # Early cache check: if file exists and size matches URL's content-length, skip re-download
     expected_size = None
     try:
@@ -241,6 +246,7 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         try:
             cached_size = output_path.stat().st_size
             if cached_size == expected_size:
+                print(f"  ✅ Cache hit: {output_path.name} ({cached_size / (1024**2):.1f} MB)")
                 # Cache hit - file exists with correct size
                 if progress_bar:
                     progress_bar(1.0)
@@ -282,6 +288,34 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
             output_path.unlink(missing_ok=True)
             raise
 
+    # Verify the downloaded file is complete
+    try:
+        actual_size = output_path.stat().st_size
+        
+        # For safetensors files, check header is valid
+        if output_path.suffix == ".safetensors":
+            import struct
+            with open(output_path, "rb") as f:
+                header_size_bytes = f.read(8)
+                if len(header_size_bytes) < 8:
+                    raise OSError(f"Safetensors file too small: {output_path.name}")
+                
+                header_size = struct.unpack("<Q", header_size_bytes)[0]
+                header = f.read(header_size)
+                if len(header) < header_size:
+                    raise OSError(f"Incomplete safetensors header in {output_path.name}")
+                
+                import json
+                json.loads(header.decode("utf-8"))
+        
+        # Verify size matches expected (if known)
+        if expected_size is not None and actual_size != expected_size:
+            print(f"  ⚠️ Size mismatch: expected {expected_size}, got {actual_size}")
+            
+    except Exception as e:
+        output_path.unlink(missing_ok=True)
+        raise OSError(f"Invalid downloaded file {output_path.name}: {str(e)}")
+
     return output_path
 
 

src/pipeline.py

@@ -78,42 +78,90 @@ def load_pipeline(
 
     try:
         set_download_cancelled(False)
+
+        # Import gr here to update button state if needed
+        import gradio as gr
+
+        print("=" * 60)
+        print("🔄 Loading SDXL Pipeline...")
+        print("=" * 60)
+
         checkpoint_filename = get_safe_filename_from_url(checkpoint_url, type_prefix="model")
         checkpoint_path = CACHE_DIR / checkpoint_filename
 
         # Check if checkpoint is already cached
         checkpoint_cached = checkpoint_path.exists() and checkpoint_path.stat().st_size > 0
 
+        # Validate cache file before using it
+        if checkpoint_cached:
+            from src.config import validate_cache_file
+            is_valid, msg = validate_cache_file(checkpoint_path)
+            if not is_valid:
+                print(f"  ⚠️ Cache invalid: {msg}")
+                checkpoint_path.unlink(missing_ok=True)
+                checkpoint_cached = False
+
         # VAE: Use suffix="_vae" and default to "vae.safetensors" for proper caching/dropdown matching
-        vae_filename = get_safe_filename_from_url(vae_url, default_name="vae.safetensors", suffix="_vae") if vae_url.strip() else "vae.safetensors"
-        vae_path = CACHE_DIR / vae_filename
-        vae_cached = vae_url.strip() and vae_path.exists() and vae_path.stat().st_size > 0
+        vae_filename = get_safe_filename_from_url(vae_url, default_name="vae.safetensors", suffix="_vae") if vae_url.strip() else None
+        vae_path = CACHE_DIR / vae_filename if vae_filename else None
+        vae_cached = vae_url.strip() and vae_path and vae_path.exists() and vae_path.stat().st_size > 0
+
+        # Validate VAE cache file before using it
+        if vae_cached:
+            from src.config import validate_cache_file
+            is_valid, msg = validate_cache_file(vae_path)
+            if not is_valid:
+                print(f"  ⚠️ VAE Cache invalid: {msg}")
+                vae_path.unlink(missing_ok=True)
+                vae_cached = False
 
         # Download checkpoint (skips if already cached)
         if progress:
             progress(0.1, desc="Downloading base model..." if not checkpoint_cached else "Loading base model...")
 
-        status_msg = f"📥 Downloading {checkpoint_path.name}..." if not checkpoint_cached else f"✅ Using cached {checkpoint_path.name}"
+        if not checkpoint_cached:
+            status_msg = f"📥 Downloading {checkpoint_path.name}..."
+            print(f"  📥 Downloading: {checkpoint_path.name}")
+        else:
+            status_msg = f"✅ Using cached {checkpoint_path.name}"
+            print(f"  ✅ Using cached: {checkpoint_path.name}")
+        
         yield status_msg, "Starting download..."
 
         if not checkpoint_cached:
             download_file_with_progress(checkpoint_url, checkpoint_path)
 
         # Download VAE if provided
-        if vae_url.strip():
-            status_msg = f"📥 Downloading {vae_path.name}..." if not vae_cached else f"✅ Using cached {vae_path.name}"
-            if progress:
-                progress(0.2, desc="Downloading VAE..." if not vae_cached else "Loading VAE...")
+        vae = None
+        if vae_url and vae_url.strip():
+            if vae_path:
+                status_msg = f"📥 Downloading {vae_path.name}..." if not vae_cached else f"✅ Using cached {vae_path.name}"
+                print(f"  📥 VAE: {vae_path.name}" if not vae_cached else f"  ✅ VAE (cached): {vae_path.name}")
+                
+                if progress:
+                    progress(0.2, desc="Downloading VAE..." if not vae_cached else "Loading VAE...")
+                
+                yield status_msg, f"Downloading VAE: {vae_path.name}" if not vae_cached else f"Using cached VAE: {vae_path.name}"
 
-            yield status_msg, f"Downloading VAE: {vae_path.name}" if not vae_cached else f"Using cached VAE: {vae_path.name}"
+                if not vae_cached:
+                    download_file_with_progress(vae_url, vae_path)
 
-            if not vae_cached:
-                download_file_with_progress(vae_url, vae_path)
+                # Load VAE from file
+                print("  ⚙️ Loading VAE weights...")
+                yield "⚙️ Loading VAE...", f"Loading VAE: {vae_path.name}"
+                vae = AutoencoderKL.from_single_file(
+                    str(vae_path),
+                    torch_dtype=dtype,
+                )
+                if progress:
+                    progress(0.25, desc="VAE loaded")
 
         # Load base pipeline (yield progress during this heavy operation)
+        print("  ⚙️ Loading SDXL pipeline from single file...")
         yield "⚙️ Loading SDXL pipeline...", "Loading model weights into memory..."
+        
         if progress:
-            progress(0.4, desc="Loading SDXL pipeline...")
+            progress(0.3, desc="Loading text encoders...")
 
         global_pipe = StableDiffusionXLPipeline.from_single_file(
             str(checkpoint_path),
@@ -121,12 +169,19 @@ def load_pipeline(
             use_safetensors=True,
             safety_checker=None,
         )
-        yield "⚙️ Pipeline loaded, setting up VAE...", f"Using device: {device_description}"
+        print("  ✅ Text encoders loaded")
+        
         if progress:
-            progress(0.6, desc="Setting up VAE...")
+            progress(0.5, desc="Loading UNet...")
 
-        if vae:
+        print("  ✅ UNet loaded")
+        yield "⚙️ Pipeline loaded, setting up components...", f"Using device: {device_description}"
+
+        # Load VAE into pipeline if provided
+        if vae is not None:
+            print("  ⚙️ Setting custom VAE...")
             global_pipe.vae = vae.to(device=device, dtype=dtype)
+            yield "⚙️ Pipeline loaded, setting up components...", f"VAE loaded: {vae_path.name}"
 
         # Parse LoRA URLs & ensure strengths list matches
         lora_urls = [u.strip() for u in lora_urls_str.split("\n") if u.strip()]
@@ -141,17 +196,29 @@ def load_pipeline(
 
         # Load and fuse each LoRA sequentially (only if URLs exist)
         if lora_urls:
+            print(f"  ⚙️ Moving pipeline to device: {device_description}...")
             global_pipe = global_pipe.to(device=device, dtype=dtype)
+
             for i, (lora_url, strength) in enumerate(zip(lora_urls, strengths)):
                 lora_filename = get_safe_filename_from_url(lora_url, suffix="_lora")
                 lora_path = CACHE_DIR / lora_filename
                 lora_cached = lora_path.exists() and lora_path.stat().st_size > 0
                 
-                status_msg = (
-                    f"📥 Downloading LoRA {i+1}/{len(lora_urls)}: {lora_path.name}..."
-                    if not lora_cached else
-                    f"✅ Using cached LoRA {i+1}/{len(lora_urls)}: {lora_path.name}"
-                )
+                # Validate LoRA cache file before using it
+                if lora_cached:
+                    from src.config import validate_cache_file
+                    is_valid, msg = validate_cache_file(lora_path)
+                    if not is_valid:
+                        print(f"  ⚠️ LoRA Cache invalid: {msg}")
+                        lora_path.unlink(missing_ok=True)
+                        lora_cached = False
+                
+                if not lora_cached:
+                    print(f"  📥 LoRA {i+1}/{len(lora_urls)}: Downloading {lora_path.name}...")
+                    status_msg = f"📥 Downloading LoRA {i+1}/{len(lora_urls)}: {lora_path.name}..."
+                else:
+                    print(f"  ✅ LoRA {i+1}/{len(lora_urls)}: Using cached {lora_path.name}")
+                    status_msg = f"✅ Using cached LoRA {i+1}/{len(lora_urls)}: {lora_path.name}"
                 
                 yield (
                     status_msg,
@@ -161,32 +228,49 @@ def load_pipeline(
                 
                 if not lora_cached:
                     download_file_with_progress(lora_url, lora_path)
-                
+
+                print(f"  ⚙️ Loading LoRA {i+1}/{len(lora_urls)}...")
                 yield f"⚙️ Loading LoRA {i+1}/{len(lora_urls)}...", f"Fusing {lora_path.name}..."
                 if progress:
                     progress(0.7 + (0.2 * i / len(lora_urls)), desc=f"Loading LoRA {i+1}/{len(lora_urls)}...")
                 
                 adapter_name = f"lora_{i}"
                 global_pipe.load_lora_weights(str(lora_path), adapter_name=adapter_name)
+                print(f"  ⚙️ Fusing LoRA {i+1} with strength={strength}...")
                 global_pipe.fuse_lora(adapter_names=[adapter_name], lora_scale=strength)
                 global_pipe.unload_lora_weights()
+        else:
+            # Move pipeline to device even without LoRAs
+            print(f"  ⚙️ Moving pipeline to device: {device_description}...")
+            global_pipe = global_pipe.to(device=device, dtype=dtype)
+
+        # Set scheduler and finalize (do this once at the end)
+        print("  ⚙️ Configuring scheduler...")
+        yield "⚙️ Finalizing pipeline...", "Setting up scheduler..."
 
-        # Set scheduler and move to device (do this once at the end)
-        yield "⚙️ Finalizing pipeline...", "Setting up scheduler and moving to device..."
         if progress:
-            progress(0.9, desc="Finalizing...")
+            progress(0.95, desc="Finalizing...")
 
-        global_pipe.scheduler.config.algorithm_type = "sde-dpmsolver++"
-        global_pipe = global_pipe.to(device=device, dtype=dtype)
+        global_pipe.scheduler = DPMSolverSDEScheduler.from_config(
+            global_pipe.scheduler.config,
+            algorithm_type="sde-dpmsolver++",
+            use_karras_sigmas=False,
+        )
 
+        print("  ✅ Pipeline ready!")
         yield "✅ Pipeline ready!", f"Ready! Loaded {len(lora_urls)} LoRA(s)"
         return ("✅ Pipeline loaded successfully!", f"Ready! Loaded {len(lora_urls)} LoRA(s)")
 
     except KeyboardInterrupt:
         set_download_cancelled(False)
+        print("\n⚠️ Download cancelled by user")
         return ("⚠️ Download cancelled by user", "Cancelled")
     except Exception as e:
-        return (f"❌ Error loading pipeline: {str(e)}", f"Error: {str(e)}")
+        import traceback
+        error_msg = f"❌ Error loading pipeline: {str(e)}"
+        print(f"\n{error_msg}")
+        print(traceback.format_exc())
+        return (error_msg, f"Error: {str(e)}")
 
 
 def cancel_download():