pearsonkyle/SDXL-Model-Merger

1"""Pipeline management for SDXL Model Merger."""2 3import torch4from diffusers import (5    StableDiffusionXLPipeline,6    AutoencoderKL,7    DPMSolverSDEScheduler,8)9 10from . import config11from .config import device, dtype, CACHE_DIR, device_description, is_running_on_spaces, set_download_cancelled12from .downloader import get_safe_filename_from_url, download_file_with_progress13from .gpu_decorator import GPU14 15 16@GPU(duration=300)17def _load_and_setup_pipeline(checkpoint_path, vae_path, lora_paths_and_strengths, load_kwargs):18    """GPU-decorated helper that performs all GPU-intensive pipeline setup."""19    _pipe = StableDiffusionXLPipeline.from_single_file(20        str(checkpoint_path),21        **load_kwargs,22    )23    print("  ✅ Text encoders loaded")24 25    # Move to device (unless using device_map='auto' which handles this automatically)26    if not is_running_on_spaces() or device != "cpu":27        print(f"  ⚙️ Moving pipeline to device: {device_description}...")28        _pipe = _pipe.to(device=device, dtype=dtype)29 30    # Load custom VAE if provided31    if vae_path is not None:32        print("  ⚙️ Loading VAE weights...")33        vae = AutoencoderKL.from_single_file(34            str(vae_path),35            torch_dtype=dtype,36        )37        print("  ⚙️ Setting custom VAE...")38        _pipe.vae = vae.to(device=device, dtype=torch.float32)39 40    # Load and fuse each LoRA41    if lora_paths_and_strengths:42        # Ensure pipeline is on device for LoRA fusion43        _pipe = _pipe.to(device=device, dtype=dtype)44 45        for i, (lora_path, strength) in enumerate(lora_paths_and_strengths):46            adapter_name = f"lora_{i}"47            print(f"  ⚙️ Loading LoRA {i+1}/{len(lora_paths_and_strengths)}...")48            _pipe.load_lora_weights(str(lora_path), adapter_name=adapter_name)49            print(f"  ⚙️ Fusing LoRA {i+1} with strength={strength}...")50            _pipe.fuse_lora(adapter_names=[adapter_name], lora_scale=strength)51            _pipe.unload_lora_weights()52    else:53        # Move pipeline to device even without LoRAs54        _pipe = _pipe.to(device=device, dtype=dtype)55 56    # Set scheduler57    print("  ⚙️ Configuring scheduler...")58    _pipe.scheduler = DPMSolverSDEScheduler.from_config(59        _pipe.scheduler.config,60        algorithm_type="sde-dpmsolver++",61        use_karras_sigmas=False,62    )63 64    # Keep VAE in float32 to prevent colorful static output65    _pipe.vae.to(dtype=torch.float32)66 67    return _pipe68 69 70def load_pipeline(71    checkpoint_url: str,72    vae_url: str,73    lora_urls_str: str,74    lora_strengths_str: str,75    progress=None76) -> tuple[str, str]:77    """78    Load SDXL pipeline with checkpoint, VAE, and LoRAs.79 80    Args:81        checkpoint_url: URL to base model .safetensors file82        vae_url: Optional URL to VAE .safetensors file83        lora_urls_str: Newline-separated URLs for LoRA models84        lora_strengths_str: Comma-separated strength values for each LoRA85        progress: Optional gr.Progress() object for UI updates86 87    Yields:88        Tuple of (status_message, progress_text) at each loading stage.89 90    Returns:91        Final yielded tuple of (final_status_message, progress_text)92    """93    # Clear any previously loaded pipeline so the UI reflects loading state94    config.set_pipe(None)95 96    try:97        set_download_cancelled(False)98 99        print("=" * 60)100        print("🔄 Loading SDXL Pipeline...")101        print("=" * 60)102 103        checkpoint_filename = get_safe_filename_from_url(checkpoint_url, type_prefix="model")104        checkpoint_path = CACHE_DIR / checkpoint_filename105 106        # Check if checkpoint is already cached107        checkpoint_cached = checkpoint_path.exists() and checkpoint_path.stat().st_size > 0108 109        # Validate cache file before using it110        if checkpoint_cached:111            is_valid, msg = config.validate_cache_file(checkpoint_path)112            if not is_valid:113                print(f"  ⚠️ Cache invalid: {msg}")114                checkpoint_path.unlink(missing_ok=True)115                checkpoint_cached = False116 117        # VAE: Use suffix="_vae" and default to "vae.safetensors" for proper caching/dropdown matching118        vae_filename = get_safe_filename_from_url(vae_url, default_name="vae.safetensors", suffix="_vae") if vae_url.strip() else None119        vae_path = CACHE_DIR / vae_filename if vae_filename else None120        vae_cached = vae_url.strip() and vae_path and vae_path.exists() and vae_path.stat().st_size > 0121 122        # Validate VAE cache file before using it123        if vae_cached:124            is_valid, msg = config.validate_cache_file(vae_path)125            if not is_valid:126                print(f"  ⚠️ VAE Cache invalid: {msg}")127                vae_path.unlink(missing_ok=True)128                vae_cached = False129 130        # Download checkpoint (skips if already cached)131        if progress:132            progress(0.1, desc="Downloading base model..." if not checkpoint_cached else "Loading base model...")133 134        if not checkpoint_cached:135            status_msg = f"📥 Downloading {checkpoint_path.name}..."136            print(f"  📥 Downloading: {checkpoint_path.name}")137        else:138            status_msg = f"✅ Using cached {checkpoint_path.name}"139            print(f"  ✅ Using cached: {checkpoint_path.name}")140 141        yield status_msg, "Starting download..."142 143        if not checkpoint_cached:144            download_file_with_progress(checkpoint_url, checkpoint_path)145 146        # Download VAE if provided (loading happens in _load_and_setup_pipeline)147        if vae_url and vae_url.strip():148            if vae_path:149                status_msg = f"📥 Downloading {vae_path.name}..." if not vae_cached else f"✅ Using cached {vae_path.name}"150                print(f"  📥 VAE: {vae_path.name}" if not vae_cached else f"  ✅ VAE (cached): {vae_path.name}")151 152                if progress:153                    progress(0.2, desc="Downloading VAE..." if not vae_cached else "Loading VAE...")154 155                yield status_msg, f"Downloading VAE: {vae_path.name}" if not vae_cached else f"Using cached VAE: {vae_path.name}"156 157                if not vae_cached:158                    download_file_with_progress(vae_url, vae_path)159 160        # For CPU/low-memory environments on Spaces, use device_map for better RAM management161        load_kwargs = {162            "torch_dtype": dtype,163            "use_safetensors": True,164        }165 166        if is_running_on_spaces() and device == "cpu":167            print("  ℹ️ CPU mode detected: enabling device_map='auto' for better RAM management")168            load_kwargs["device_map"] = "auto"169 170        # Parse LoRA URLs & ensure strengths list matches171        lora_urls = [u.strip() for u in lora_urls_str.split("\n") if u.strip()]172        strengths_raw = [s.strip() for s in lora_strengths_str.split(",")]173        strengths = []174        for i, url in enumerate(lora_urls):175            try:176                val = float(strengths_raw[i]) if i < len(strengths_raw) else 1.0177                strengths.append(val)178            except ValueError:179                strengths.append(1.0)180 181        # Download LoRAs (CPU-bound downloads, before GPU work)182        lora_paths_and_strengths = []183        if lora_urls:184            for i, (lora_url, strength) in enumerate(zip(lora_urls, strengths)):185                lora_filename = get_safe_filename_from_url(lora_url, suffix="_lora")186                lora_path = CACHE_DIR / lora_filename187                lora_cached = lora_path.exists() and lora_path.stat().st_size > 0188 189                # Validate LoRA cache file before using it190                if lora_cached:191                    is_valid, msg = config.validate_cache_file(lora_path)192                    if not is_valid:193                        print(f"  ⚠️ LoRA Cache invalid: {msg}")194                        lora_path.unlink(missing_ok=True)195                        lora_cached = False196 197                if not lora_cached:198                    print(f"  📥 LoRA {i+1}/{len(lora_urls)}: Downloading {lora_path.name}...")199                    status_msg = f"📥 Downloading LoRA {i+1}/{len(lora_urls)}: {lora_path.name}..."200                else:201                    print(f"  ✅ LoRA {i+1}/{len(lora_urls)}: Using cached {lora_path.name}")202                    status_msg = f"✅ Using cached LoRA {i+1}/{len(lora_urls)}: {lora_path.name}"203 204                yield (205                    status_msg,206                    f"Downloading LoRA {i+1}/{len(lora_urls)} ({lora_path.name})..." if not lora_cached207                    else f"Using cached LoRA {i+1}/{len(lora_urls)} ({lora_path.name})"208                )209 210                if not lora_cached:211                    download_file_with_progress(lora_url, lora_path)212 213                lora_paths_and_strengths.append((lora_path, strength))214 215        # All downloads complete — now do GPU-intensive setup in one decorated call216        yield "⚙️ Loading SDXL pipeline...", "Loading model weights into memory..."217 218        if progress:219            progress(0.5, desc="Loading pipeline...")220 221        _pipe = _load_and_setup_pipeline(222            checkpoint_path, vae_path, lora_paths_and_strengths, load_kwargs223        )224 225        if progress:226            progress(0.95, desc="Finalizing...")227 228        # ✅ Only publish the pipeline globally AFTER all steps succeed229        config.set_pipe(_pipe)230 231        print("  ✅ Pipeline ready!")232        yield "✅ Pipeline ready!", f"Ready! Loaded {len(lora_urls)} LoRA(s)"233 234    except KeyboardInterrupt:235        set_download_cancelled(False)236        config.set_pipe(None)237        print("\n⚠️ Download cancelled by user")238        return ("⚠️ Download cancelled by user", "Cancelled")239    except Exception as e:240        import traceback241        config.set_pipe(None)242        error_msg = f"❌ Error loading pipeline: {str(e)}"243        print(f"\n{error_msg}")244        print(traceback.format_exc())245        return (error_msg, f"Error: {str(e)}")246 247 248def cancel_download():249    """Set the global cancellation flag to stop any ongoing downloads."""250    set_download_cancelled(True)251 252 253def get_pipeline() -> StableDiffusionXLPipeline | None:254    """Get the currently loaded pipeline."""255    return config.get_pipe()256