pearsonkyle/SDXL-Model-Merger

1"""Download utilities for SDXL Model Merger with Gradio progress integration."""2 3import re4import requests5from pathlib import Path6from tqdm import tqdm as TqdmBase7 8from .config import download_cancelled9 10 11def extract_model_id(url: str) -> str | None:12    """Extract CivitAI model ID from URL."""13    match = re.search(r'/models/(\d+)', url)14    return match.group(1) if match else None15 16 17def is_huggingface_url(url: str) -> bool:18    """Check if URL is a HuggingFace model download URL."""19    return "huggingface.co" in url.lower()20 21 22def get_safe_filename_from_url(23    url: str,24    default_name: str = "model.safetensors",25    suffix: str = "",26    type_prefix: str | None = None27) -> str:28    """29    Generate a safe filename with model ID from URL.30 31    For CivitAI URLs like https://civitai.com/api/download/models/12345?type=...32 33    Naming patterns:34    - Checkpoint (type_prefix='model'): 12345_model.safetensors or 12345_model_anime_style.safetensors35    - VAE (suffix='_vae'): 12345_vae.safetensors (no name extraction to avoid double suffix)36    - LoRA (suffix='_lora'): 12345_lora.safetensors (no name extraction to avoid double suffix)37 38    For HuggingFace URLs without model IDs, attempts to extract name from path or uses suffix-based naming.39 40    Args:41        url: The download URL42        default_name: Fallback filename if extraction fails43        suffix: Optional suffix to append before .safetensors (e.g., '_vae', '_lora')44        type_prefix: Optional prefix after model_id (e.g., 'model' -> 12345_model.safetensors)45    """46    model_id = extract_model_id(url)47 48    # If no CivitAI model ID, try to generate a name from HuggingFace path49    if not model_id and "huggingface.co" in url:50        # Try to extract name from URL path (e.g., sdxl-vae-fp16-fix -> fp16_fix)51        try:52            parts = url.split("huggingface.co/")[1] if "huggingface.co/" in url else ""53            if parts:54                # Get the repo name (second part after org/)55                path_parts = [p for p in parts.split("/") if p]56                if len(path_parts) >= 2:57                    repo_name = path_parts[1]58                    # Clean up and create a simple identifier59                    clean_repo = re.sub(r'[^a-zA-Z0-9]', '_', repo_name)[:30].strip('_')60                    if clean_repo:61                        model_id = f"hf_{clean_repo}"62        except Exception:63            pass64 65    if not model_id:66        return default_name67 68    # Special handling for VAE/LoRA with HuggingFace URLs to avoid double suffix69    is_special_type = suffix in ("_vae", "_lora")70    71    # Strip common suffixes from model_id when adding corresponding suffix72    # (e.g., "sdxl_vae_fp16_fix" + "_vae" -> "sdxl_fp16_fix" + "_vae")73    if is_special_type:74        strip_suffix = suffix.lstrip('_')  # "vae" or "lora"75        model_id_lower = model_id.lower()76        # Check if model_id contains the type (with underscore boundaries)77        if f"_{strip_suffix}_" in model_id_lower or model_id_lower.endswith(f"_{strip_suffix}"):78            # Remove the suffix from model_id79            if model_id_lower.endswith(f"_{strip_suffix}"):80                model_id = model_id[:-len(strip_suffix)-1]81            else:82                # Find and remove _suffix_ pattern83                pattern = f"_{strip_suffix}_"84                idx = model_id_lower.find(pattern)85                if idx >= 0:86                    model_id = model_id[:idx] + model_id[idx+len(pattern):]87 88    # Build the name portion: either clean name from URL or fallback89    name_part = ""90 91    # For VAE/LoRA types, skip Content-Disposition parsing to avoid double naming92    # (e.g., sdxl_vae_vae instead of just vae)93    if not is_special_type:94        try:95            response = requests.head(url, timeout=10, allow_redirects=True)96            cd = response.headers.get('Content-Disposition', '')97            match = re.search(r'filename="([^"]+)"', cd)98            if match:99                filename = match.group(1)100                # Extract base name without extension101                base_name = Path(filename).stem102                # Clean up the name (remove special chars)103                clean_name = re.sub(r'[^\w\s-]', '', base_name)[:50]104                clean_name = re.sub(r'[-\s]+', '_', clean_name.strip('-_'))105                if clean_name:106                    name_part = clean_name107        except Exception:108            pass109 110    # Build filename with model_id, optional type_prefix, optional name_part, and suffix111    parts = [model_id]112    if type_prefix:113        parts.append(type_prefix)114    if name_part:115        parts.append(name_part)116 117    # Handle suffix - for VAE/LoRA we only add the suffix, not double naming118    if suffix:119        if is_special_type:120            # For _vae and _lora: just use model_id + suffix directly121            return f"{model_id}{suffix}.safetensors"122        else:123            # For other types (checkpoint), append suffix after name_part124            parts.append(suffix.lstrip('_'))125 126    return '_'.join(p for p in parts if p).replace('__', '_') + '.safetensors'127 128 129class TqdmGradio(TqdmBase):130    """tqdm subclass that sends progress updates to Gradio's gr.Progress()"""131 132    def __init__(self, *args, gradio_prog=None, **kwargs):133        super().__init__(*args, **kwargs)134        self.gradio_prog = gradio_prog135        self.last_pct = 0136 137    def update(self, n=1):138        from .config import download_cancelled139        if download_cancelled:140            raise KeyboardInterrupt("Download cancelled by user")141        super().update(n)142        if self.gradio_prog and self.total:143            pct = int(100 * self.n / self.total)144            # Only update UI every ~5% to avoid spamming145            if pct != self.last_pct and pct % 5 == 0:146                self.last_pct = pct147                self.gradio_prog(pct / 100)148 149 150def get_cached_file_size(url: str, suffix: str = "", type_prefix: str | None = None) -> tuple[Path | None, int | None]:151    """152    Check if file exists in cache and matches expected size.153 154    Uses the same filename generation logic as download operations to find155    cached files by URL.156 157    Args:158        url: The download URL to check for cached file159        suffix: Optional suffix (e.g., '_vae', '_lora') for special file types160        type_prefix: Optional prefix after model_id (e.g., 'model')161 162    Returns:163        Tuple of (cached_file_path, file_size) if valid cache exists,164        or (None, None) if no valid cache found.165    """166    from .config import CACHE_DIR167 168    # Generate the expected filename for this URL169    default_name = "vae.safetensors" if suffix == "_vae" else (170        "lora.safetensors" if suffix == "_lora" else "model.safetensors"171    )172 173    cached_filename = get_safe_filename_from_url(174        url,175        default_name=default_name,176        suffix=suffix,177        type_prefix=type_prefix178    )179 180    cached_path = CACHE_DIR / cached_filename181 182    if cached_path.exists() and cached_path.is_file():183        try:184            file_size = cached_path.stat().st_size185            # Only return valid cache if file has content186            if file_size > 0:187                return cached_path, file_size188        except OSError:189            pass190 191    return None, None192 193 194def download_file_with_progress(url: str, output_path: Path, progress_bar=None) -> Path:195    """196    Download a file with Gradio-synced progress bar + cancel support.197 198    Checks for existing cached files before downloading. If a valid cache199    exists (file exists with matching expected size), skips re-download.200    201    Supports both HTTP(S) and HuggingFace Hub URLs.202 203    Args:204        url: File URL to download (http/https/file)205        output_path: Destination path for downloaded file206        progress_bar: Optional gr.Progress() object for UI updates207 208    Returns:209        Path to the downloaded (or cached) file210 211    Raises:212        KeyboardInterrupt: If download is cancelled213        requests.RequestException: If download fails214    """215    from .config import download_cancelled216 217    # Handle local file:// URLs218    if url.startswith("file://"):219        local_path = Path(url[7:])  # Remove "file://" prefix220        if local_path.exists():221            import shutil222            output_path.parent.mkdir(parents=True, exist_ok=True)223 224            print(f"  📁 Copying from cache: {local_path.name} → {output_path.name}")225 226            # Copy the file to cache location227            shutil.copy2(str(local_path), str(output_path))228 229            # Update progress bar for cached files230            if progress_bar:231                progress_bar(1.0)232            return output_path233        else:234            raise FileNotFoundError(f"Local file not found: {local_path}")235 236    print(f"  📥 Downloading to cache: {output_path.name}")237    238    # Early cache check: if file exists and size matches URL's content-length, skip re-download239    expected_size = None240    try:241        head = requests.head(url, timeout=10)242        expected_size = int(head.headers.get('content-length', 0))243    except Exception:244        pass  # Skip header fetch on errors245 246    if output_path.exists() and expected_size is not None:247        try:248            cached_size = output_path.stat().st_size249            if cached_size == expected_size:250                print(f"  ✅ Cache hit: {output_path.name} ({cached_size / (1024**2):.1f} MB)")251                # Cache hit - file exists with correct size252                if progress_bar:253                    progress_bar(1.0)254                return output_path  # Skip re-download!255        except OSError:256            pass  # File access error, proceed with download257 258    output_path.parent.mkdir(parents=True, exist_ok=True)259 260    session = requests.Session()261    response = session.get(url, stream=True, timeout=30)262    response.raise_for_status()263 264    total_size = expected_size or int(response.headers.get('content-length', 0))265    block_size = 8192266 267    # Use TqdmGradio to sync progress with Gradio268    tqdm_kwargs = {269        'unit': 'B',270        'unit_scale': True,271        'desc': f"Downloading {output_path.name}",272        'gradio_prog': progress_bar,273        'disable': False,274        'bar_format': '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]',275    }276 277    with open(output_path, "wb") as f:278        try:279            for data in TqdmGradio(280                response.iter_content(block_size),281                total=total_size // block_size if total_size else 0,282                **tqdm_kwargs,283            ):284                if download_cancelled:285                    raise KeyboardInterrupt("Download cancelled by user")286                f.write(data)287        except KeyboardInterrupt:288            # Clean partial file on cancel289            output_path.unlink(missing_ok=True)290            raise291 292    # Verify the downloaded file is complete293    try:294        actual_size = output_path.stat().st_size295        296        # For safetensors files, check header is valid297        if output_path.suffix == ".safetensors":298            import struct299            with open(output_path, "rb") as f:300                header_size_bytes = f.read(8)301                if len(header_size_bytes) < 8:302                    raise OSError(f"Safetensors file too small: {output_path.name}")303                304                header_size = struct.unpack("<Q", header_size_bytes)[0]305                header = f.read(header_size)306                if len(header) < header_size:307                    raise OSError(f"Incomplete safetensors header in {output_path.name}")308                309                import json310                json.loads(header.decode("utf-8"))311        312        # Verify size matches expected (if known)313        if expected_size is not None and actual_size != expected_size:314            print(f"  ⚠️ Size mismatch: expected {expected_size}, got {actual_size}")315            316    except Exception as e:317        output_path.unlink(missing_ok=True)318        raise OSError(f"Invalid downloaded file {output_path.name}: {str(e)}")319 320    return output_path321 322 323def clear_cache(cache_dir: Path = None, keep_extensions: list[str] = None):324    """325    Remove old cache files.326 327    Args:328        cache_dir: Cache directory path (defaults to config.CACHE_DIR)329        keep_extensions: File extensions to preserve (default: ['.safetensors'])330    """331    if cache_dir is None:332        from .config import CACHE_DIR333        cache_dir = CACHE_DIR334 335    if keep_extensions is None:336        keep_extensions = ['.safetensors']337 338    # Remove temp files339    for file in cache_dir.glob("*.tmp"):340        file.unlink()341 342    # Optional: age-based cleanup (7 days)343    # import time344    # cutoff = time.time() - 86400 * 7345    # for f in cache_dir.iterdir():346    #     if f.is_file() and f.stat().st_mtime < cutoff:347    #         f.unlink()348