Update Gradio to 6.9.0, add cached file checks in LoRA loading, fix unloading errors, optimize download cancellation, improve device/dtype handling, update optimum dependency.

app.py

@@ -83,7 +83,7 @@ def create_app():
     from src.exporter import export_merged_model
     from src.config import get_cached_models, get_cached_checkpoints, get_cached_vaes, get_cached_loras
 
-    with gr.Blocks(title="SDXL Model Merger", css=header_css) as demo:
+    with gr.Blocks(title="SDXL Model Merger") as demo:
         # Header section
         with gr.Column(elem_classes=["feature-card"]):
             gr.HTML("""
@@ -393,7 +393,18 @@ def create_app():
             fn=load_pipeline,
             inputs=[checkpoint_url, vae_url, lora_urls, lora_strengths],
             outputs=[load_status, load_progress],
-            show_api=False,
+        ).then(
+            fn=on_load_pipeline_complete,
+            inputs=[load_status, load_progress],
+            outputs=[load_status, load_progress, load_btn],
+        ).then(
+            fn=lambda: (
+                gr.update(choices=["(None found)"] + get_cached_checkpoints()),
+                gr.update(choices=["(None found)"] + get_cached_vaes()),
+                gr.update(choices=["(None found)"] + get_cached_loras()),
+            ),
+            inputs=[],
+            outputs=[cached_checkpoints, cached_vaes, cached_loras],
         )
 
         def on_cached_checkpoint_change(cached_path):
@@ -423,10 +434,10 @@ def create_app():
         def on_cached_lora_change(cached_path, current_urls):
             """Add cached LoRA to the list."""
             if cached_path and cached_path != "(None found)":
-                # Add new LoRA to existing URLs (avoid duplicate)
                 urls_list = [u.strip() for u in current_urls.split("\n") if u.strip()]
-                if cached_path not in urls_list:
-                    urls_list.append(cached_path)
+                file_url = f"file://{cached_path}"
+                if file_url not in urls_list:
+                    urls_list.append(file_url)
                     return gr.update(value="\n".join(urls_list))
             return gr.update()
 
@@ -470,7 +481,6 @@ def create_app():
             fn=generate_image,
             inputs=[prompt, negative_prompt, cfg, steps, height, width, tile_x, tile_y],
             outputs=[image_output, gen_progress],
-            show_api=False,
         ).then(
             fn=lambda img, msg: on_generate_complete(msg, "Done", img),
             inputs=[image_output, gen_progress],
@@ -515,7 +525,6 @@ def create_app():
             ),
             inputs=[include_lora, quantize_toggle, qtype_dropdown, format_dropdown],
             outputs=[download_link, export_progress],
-            show_api=False,
         ).then(
             fn=lambda path, msg: on_export_complete(msg, "Exported", path),
             inputs=[download_link, export_progress],
@@ -534,7 +543,7 @@ def create_app():
 def main():
     """Create and launch the Gradio app."""
     app = create_app()
-    # CSS is embedded in the Blocks, so we pass it to launch for Gradio 6+
+    # CSS is passed to launch() in Gradio 6+
     app.launch()
 
 

requirements.txt

@@ -5,12 +5,13 @@ torch>=2.0.0
 diffusers>=0.24.0
 transformers>=4.35.0
 safetensors>=0.4.0
+optimum>=1.0.0
 
 # Image processing
 Pillow>=10.0.0
 
 # UI framework
-gradio>=4.0.0
+gradio>=6.9.0
 
 # Download utilities
 tqdm>=4.65.0

src/downloader.py

@@ -127,18 +127,56 @@ def set_download_cancelled(value: bool):
     download_cancelled = value
 
 
-def get_cached_file_size(url: str) -> tuple[Path | None, int | None]:
+def get_cached_file_size(url: str, suffix: str = "", type_prefix: str | None = None) -> tuple[Path | None, int | None]:
     """
     Check if file exists in cache and matches expected size.
-    Returns (path, expected_size) or (None, None) if no valid cache.
+    
+    Uses the same filename generation logic as download operations to find
+    cached files by URL.
+
+    Args:
+        url: The download URL to check for cached file
+        suffix: Optional suffix (e.g., '_vae', '_lora') for special file types
+        type_prefix: Optional prefix after model_id (e.g., 'model')
+
+    Returns:
+        Tuple of (cached_file_path, file_size) if valid cache exists,
+        or (None, None) if no valid cache found.
     """
-    # Simple implementation - would need URL-to-filename mapping for production
+    from .config import CACHE_DIR
+    
+    # Generate the expected filename for this URL
+    default_name = "vae.safetensors" if suffix == "_vae" else (
+        "lora.safetensors" if suffix == "_lora" else "model.safetensors"
+    )
+    
+    cached_filename = get_safe_filename_from_url(
+        url, 
+        default_name=default_name,
+        suffix=suffix,
+        type_prefix=type_prefix
+    )
+    
+    cached_path = CACHE_DIR / cached_filename
+    
+    if cached_path.exists() and cached_path.is_file():
+        try:
+            file_size = cached_path.stat().st_size
+            # Only return valid cache if file has content
+            if file_size > 0:
+                return cached_path, file_size
+        except OSError:
+            pass
+    
     return None, None
 
 
 def download_file_with_progress(url: str, output_path: Path, progress_bar=None) -> Path:
     """
     Download a file with Gradio-synced progress bar + cancel support.
+    
+    Checks for existing cached files before downloading. If a valid cache
+    exists (file exists with matching expected size), skips re-download.
 
     Args:
         url: File URL to download (http/https/file)
@@ -146,14 +184,13 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         progress_bar: Optional gr.Progress() object for UI updates
 
     Returns:
-        Path to the downloaded file
+        Path to the downloaded (or cached) file
 
     Raises:
         KeyboardInterrupt: If download is cancelled
         requests.RequestException: If download fails
     """
     global download_cancelled
-    download_cancelled = False
 
     # Handle local file:// URLs
     if url.startswith("file://"):
@@ -163,7 +200,7 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
             output_path.parent.mkdir(parents=True, exist_ok=True)
             # Copy the file to cache location
             shutil.copy2(str(local_path), str(output_path))
-            
+
             # Update progress bar for cached files
             if progress_bar:
                 progress_bar(1.0)
@@ -171,18 +208,24 @@ def download_file_with_progress(url: str, output_path: Path, progress_bar=None)
         else:
             raise FileNotFoundError(f"Local file not found: {local_path}")
 
-    # Cache check: if file exists and size matches URL's content-length, skip re-download
+    # Early cache check: if file exists and size matches URL's content-length, skip re-download
     expected_size = None
     try:
         head = requests.head(url, timeout=10)
         expected_size = int(head.headers.get('content-length', 0))
-        if output_path.exists() and output_path.stat().st_size == expected_size:
-            # Cache hit - still update progress to show completion
-            if progress_bar:
-                progress_bar(1.0)
-            return output_path  # Cache hit!
     except Exception:
-        pass  # Skip cache validation on errors
+        pass  # Skip header fetch on errors
+    
+    if output_path.exists() and expected_size is not None:
+        try:
+            cached_size = output_path.stat().st_size
+            if cached_size == expected_size:
+                # Cache hit - file exists with correct size
+                if progress_bar:
+                    progress_bar(1.0)
+                return output_path  # Skip re-download!
+        except OSError:
+            pass  # File access error, proceed with download
 
     output_path.parent.mkdir(parents=True, exist_ok=True)
 

src/exporter.py

@@ -34,7 +34,10 @@ def export_merged_model(
         # Step 1: Unload LoRAs
         yield "💾 Exporting model...", "Unloading LoRAs..."
         if include_lora:
-            global_pipe.unload_lora_weights()
+            try:
+                global_pipe.unload_lora_weights()
+            except Exception:
+                pass
 
         merged_state_dict = {}
 

src/pipeline.py

@@ -9,8 +9,8 @@ from diffusers import (
     DPMSolverSDEScheduler,
 )
 
-from .config import device, dtype, pipe as global_pipe, CACHE_DIR, download_cancelled, device_description
-from .downloader import download_file_with_progress, get_safe_filename_from_url
+from .config import device, dtype, pipe as global_pipe, CACHE_DIR, device_description
+from .downloader import download_file_with_progress, get_safe_filename_from_url, set_download_cancelled
 
 
 def _make_asymmetric_forward(module, pad_h: int, pad_w: int, tile_x: bool, tile_y: bool):
@@ -74,28 +74,41 @@ def load_pipeline(
     Returns:
         Tuple of (final_status_message, progress_text)
     """
-    global global_pipe, download_cancelled
+    global global_pipe
 
     try:
+        set_download_cancelled(False)
         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
         
         # 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
 
-        # Download checkpoint
+        # Download checkpoint (skips if already cached)
         if progress:
-            progress(0.1, desc="Downloading base model...")
-        yield f"📥 Downloading {checkpoint_path.name}...", "Starting download..."
-        download_file_with_progress(checkpoint_url, checkpoint_path)
+            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}"
+        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...")
-            yield f"📥 Downloading {vae_path.name}...", f"Downloading VAE: {vae_path.name}"
-            download_file_with_progress(vae_url, vae_path)
+                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}"
+            
+            if not vae_cached:
+                download_file_with_progress(vae_url, vae_path)
             vae = AutoencoderKL.from_single_file(str(vae_path), torch_dtype=dtype)
         else:
             vae = None
@@ -126,27 +139,30 @@ def load_pipeline(
 
         # Load and fuse each LoRA sequentially (only if URLs exist)
         if lora_urls:
-            first_lora_filename = get_safe_filename_from_url(lora_urls[0], "lora_0.safetensors", suffix="_lora")
-            first_lora_path = CACHE_DIR / first_lora_filename
-            yield f"📥 Downloading LoRA: {first_lora_path.name}...", f"Downloading LoRA 1/... ({first_lora_path.name})..."
-            download_file_with_progress(lora_urls[0], first_lora_path)
-
-            global_pipe.load_lora_weights(str(first_lora_path), adapter_name="main_lora")
-            global_pipe.fuse_lora(adapter_names=["main_lora"], lora_scale=strengths[0])
-
-            for i in range(1, len(lora_urls)):
-                lora_filename = get_safe_filename_from_url(lora_urls[i], f"lora_{i}.safetensors", suffix="_lora")
+            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, f"lora_{i}.safetensors", suffix="_lora")
                 lora_path = CACHE_DIR / lora_filename
-                yield f"📥 Downloading LoRA {i+1}...", f"Downloading LoRA {i+1}/{len(lora_urls)} ({lora_path.name})..."
-                download_file_with_progress(lora_urls[i], lora_path)
-
-                global_pipe.unload_lora_weights()
-                global_pipe.load_lora_weights(str(lora_path), adapter_name=f"lora_{i}")
-                # Fuse all loaded adapters so far
-                global_pipe.fuse_lora(
-                    adapter_names=["main_lora"] + [f"lora_{j}" for j in range(1, i+1)],
-                    lora_scale=strengths[i]
+                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}"
                 )
+                
+                yield (
+                    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})"
+                )
+                
+                if not lora_cached:
+                    download_file_with_progress(lora_url, lora_path)
+                adapter_name = f"lora_{i}"
+                global_pipe.load_lora_weights(str(lora_path), adapter_name=adapter_name)
+                global_pipe.fuse_lora(adapter_names=[adapter_name], lora_scale=strength)
+                global_pipe.unload_lora_weights()
 
         # Set scheduler and move to device (do this once at the end)
         yield "⚙️ Finalizing...", "Setting up scheduler..."
@@ -157,7 +173,7 @@ def load_pipeline(
         return ("✅ Pipeline loaded successfully!", f"Ready! Loaded {len(lora_urls)} LoRA(s)")
 
     except KeyboardInterrupt:
-        download_cancelled = False
+        set_download_cancelled(False)
         return ("⚠️ Download cancelled by user", "Cancelled")
     except Exception as e:
         return (f"❌ Error loading pipeline: {str(e)}", f"Error: {str(e)}")
@@ -165,8 +181,7 @@ def load_pipeline(
 
 def cancel_download():
     """Set the global cancellation flag to stop any ongoing downloads."""
-    global download_cancelled
-    download_cancelled = True
+    set_download_cancelled(True)
 
 
 def get_pipeline() -> StableDiffusionXLPipeline | None: