from huggingface_hub import HfApi, list_repo_files, hf_hub_url, get_hf_file_metadata import os import argparse import requests import threading import time import sys import hashlib import textwrap from pathlib import Path import concurrent.futures from typing import List, Dict, Optional, Tuple import shutil import json # Configuration DEFAULT_DIRS = { "qwen": "Training_Models_Qwen", "wan": "Training_Models_Wan", "flux2_klein": "Training_Models_FLUX_2_and_Klein", "z_image": "Training_Models_Z_Image", } # Model configurations MODEL_CONFIGS = { "qwen_image": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "qwen_2.5_vl_7b_bf16.safetensors", "qwen_train_vae.safetensors", "qwen_image_bf16.safetensors" ], "name": "Qwen Image Training Models", "description": "Qwen 2.5 VL 7B model with BF16 precision for image training", "default_dir": "Training_Models_Qwen" }, "qwen_image_2512": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "qwen_2.5_vl_7b_bf16.safetensors", "qwen_train_vae.safetensors", "Qwen_Image_2512_BF16.safetensors" ], "name": "Qwen Image (2512) Training Models", "description": "Qwen Image 2512 model with BF16 precision for image training", "default_dir": "Training_Models_Qwen" }, "qwen_image_edit_plus": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "qwen_2.5_vl_7b_bf16.safetensors", "qwen_train_vae.safetensors", "Qwen_Image_Edit_Plus_2509_bf16.safetensors" ], "name": "Qwen Image Edit Plus (2509) Training Models", "description": "Qwen Image Edit Plus 2509 model with BF16 precision for image editing training with multiple control images", "default_dir": "Training_Models_Qwen" }, "qwen_image_edit_2511": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "qwen_2.5_vl_7b_bf16.safetensors", "qwen_train_vae.safetensors", "Qwen_Image_Edit_2511_BF16.safetensors" ], "name": "Qwen Image Edit (2511) Training Models", "description": "Qwen Image Edit 2511 model with BF16 precision for image editing training with multiple control images", "default_dir": "Training_Models_Qwen" }, "wan21_t2v": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "Wan2_1_VAE_bf16.safetensors", "models_clip_open-clip-xlm-roberta-large-vit-huge-14.safetensors", "umt5-xxl-enc-bf16.safetensors", "wan2.1_t2v_14B_bf16.safetensors" ], "name": "Wan 2.1 Text to Video Training Models", "description": "Wan 2.1 14B model with BF16 precision for text-to-video training", "default_dir": "Training_Models_Wan" }, "wan22_t2v": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "Wan2_1_VAE_bf16.safetensors", "umt5-xxl-enc-bf16.safetensors", "Wan-2.2-T2V-High-Noise-BF16.safetensors", "Wan-2.2-T2V-Low-Noise-BF16.safetensors" ], "name": "Wan 2.2 Text to Video Training Models", "description": "Wan 2.2 14B models with BF16 precision for text-to-video training (high and low noise variants)", "default_dir": "Training_Models_Wan" }, "wan22_i2v": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ "Wan2_1_VAE_bf16.safetensors", "umt5-xxl-enc-bf16.safetensors", "Wan-2.2-I2V-Low-Noise-BF16.safetensors", "Wan-2.2-I2V-High-Noise-BF16.safetensors" ], "name": "Wan 2.2 Image to Video Training Models", "description": "Wan 2.2 14B models with BF16 precision for image-to-video training (high and low noise variants)", "default_dir": "Training_Models_Wan" }, "flux2_dev": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ {"remote": "FLUX_2_Models/FLUX_2_Dev_BF16.safetensors", "local": "FLUX_2_Dev_BF16.safetensors"}, {"remote": "FLUX_2_Models/FLUX_2_Klein_Train_VAE.safetensors", "local": "FLUX_2_Klein_Train_VAE.safetensors"}, {"remote": "FLUX_2_Models/Mistral3_FLUX2_BF16.safetensors", "local": "Mistral3_FLUX2_BF16.safetensors"}, ], "name": "FLUX 2 Dev Training Models", "description": "FLUX 2 Dev bundle (BF16) with Klein training VAE and Mistral3 text encoder", "default_dir": "Training_Models_FLUX_2_and_Klein", }, "flux_klein_9b": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ {"remote": "FLUX_2_Models/FLUX2-Klein-Base-9B.safetensors", "local": "FLUX2-Klein-Base-9B.safetensors"}, {"remote": "FLUX_2_Models/FLUX_2_Klein_Train_VAE.safetensors", "local": "FLUX_2_Klein_Train_VAE.safetensors"}, {"remote": "FLUX_2_Models/qwen_3_8b.safetensors", "local": "qwen_3_8b.safetensors"}, ], "name": "FLUX Klein 9B Training Models", "description": "FLUX Klein Base 9B bundle with Klein training VAE and qwen_3_8b", "default_dir": "Training_Models_FLUX_2_and_Klein", }, "flux_klein_4b": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ {"remote": "FLUX_2_Models/FLUX2-Klein-Base-4B.safetensors", "local": "FLUX2-Klein-Base-4B.safetensors"}, {"remote": "FLUX_2_Models/FLUX_2_Klein_Train_VAE.safetensors", "local": "FLUX_2_Klein_Train_VAE.safetensors"}, {"remote": "FLUX_2_Models/qwen_3_4b.safetensors", "local": "qwen_3_4b.safetensors"}, ], "name": "FLUX Klein 4B Training Models", "description": "FLUX Klein Base 4B bundle with Klein training VAE and qwen_3_4b", "default_dir": "Training_Models_FLUX_2_and_Klein", }, "z_image_base": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ {"remote": "FLUX_2_Models/Z_Image_Train_VAE.safetensors", "local": "Z_Image_Train_VAE.safetensors"}, {"remote": "Z_Image_Training_Text_Encoder.safetensors", "local": "Z_Image_Training_Text_Encoder.safetensors"}, {"remote": "Z_Image_BF16.safetensors", "local": "Z_Image_BF16.safetensors"}, ], "name": "Z-Image Base Training Models", "description": "Z-Image Base bundle (BF16) with Z_Image training VAE and Z_Image_Training_Text_Encoder", "default_dir": "Training_Models_Z_Image", }, "z_image_turbo": { "repo_id": "MonsterMMORPG/Wan_GGUF", "files": [ {"remote": "FLUX_2_Models/Z_Image_Train_VAE.safetensors", "local": "Z_Image_Train_VAE.safetensors"}, {"remote": "Z_Image_Training_Text_Encoder.safetensors", "local": "Z_Image_Training_Text_Encoder.safetensors"}, {"remote": "FLUX_2_Models/zimage_turbo_training_adapter_v2.safetensors", "local": "zimage_turbo_training_adapter_v2.safetensors"}, {"remote": "Z-Image-Turbo-Models/Z_Image_Turbo_BF16.safetensors", "local": "Z_Image_Turbo_BF16.safetensors"}, ], "name": "Z-Image Turbo Training Models", "description": "Z-Image Turbo bundle (BF16) with training VAE, Z_Image_Training_Text_Encoder, and turbo adapter v2", "default_dir": "Training_Models_Z_Image", } } # Interactive menu order (must stay in sync with Windows_Download_Training_Model_Files.bat) MODEL_MENU = [ "qwen_image", "qwen_image_2512", "qwen_image_edit_plus", "qwen_image_edit_2511", "wan21_t2v", "wan22_t2v", "wan22_i2v", "flux2_dev", "flux_klein_9b", "flux_klein_4b", "z_image_base", "z_image_turbo", ] DOWNLOAD_CONFIG = { "num_connections": 16, # Fixed 16 connections "chunk_size": 10485760, # 10MB buffer for streaming "max_retries": 5, # Reasonable retry count "retry_delay": 2, # Base delay between retries "max_retry_delay": 30, # Cap exponential backoff "timeout": 300, # 5 minute timeout per request } class RobustDownloader: def __init__(self, config: Dict): self.config = config # Create session with connection pooling self.session = requests.Session() adapter = requests.adapters.HTTPAdapter( pool_connections=20, pool_maxsize=20, max_retries=3 ) self.session.mount('http://', adapter) self.session.mount('https://', adapter) self.session.headers.update({ 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36' }) # Cache files in the same directory where script runs script_dir = os.path.dirname(os.path.abspath(__file__)) if __file__ else os.getcwd() # SHA256 cache file self.sha_cache_file = os.path.join(script_dir, "sha256_cache.json") self.sha_cache = self.load_sha_cache() # Verified files cache to avoid re-verification self.verified_cache_file = os.path.join(script_dir, "verified_files_cache.json") self.verified_cache = self.load_verified_cache() # Debug: Print cache file locations print(f"[DEBUG] Cache files will be saved in: {script_dir}") print(f"[DEBUG] SHA256 cache: {self.sha_cache_file}") print(f"[DEBUG] Verified cache: {self.verified_cache_file}") print(f"[DEBUG] Current working directory: {os.getcwd()}") # Console/progress management (single-line, cross-platform) self._progress_lock = threading.Lock() self._active_progress = False self._last_progress_len = 0 # ------------- Console helpers to ensure single-line progress ------------- def _get_terminal_width(self) -> int: try: return shutil.get_terminal_size(fallback=(100, 20)).columns except Exception: return 100 def _shorten_middle(self, text: str, max_len: int) -> str: if len(text) <= max_len: return text if max_len <= 3: return text[:max_len] left = (max_len - 3) // 2 right = (max_len - 3) - left return text[:left] + "..." + text[-right:] def _clear_progress_line_locked(self): # Clear current line safely width = self._get_terminal_width() clear_len = max(self._last_progress_len, width) sys.stdout.write("\r" + " " * clear_len + "\r") sys.stdout.flush() self._last_progress_len = 0 self._active_progress = False def clear_progress_line(self): with self._progress_lock: if self._active_progress: self._clear_progress_line_locked() def show_progress_line(self, text: str): # Ensure we never wrap: truncate to terminal width - 1 with self._progress_lock: width = self._get_terminal_width() max_len = max(1, width - 1) if len(text) > max_len: text = text[:max_len] # Write text and erase any remainder from previous longer line sys.stdout.write("\r" + text) extra_spaces = max(0, max(self._last_progress_len - len(text), 0)) if extra_spaces: sys.stdout.write(" " * extra_spaces) sys.stdout.write("\r" + text) sys.stdout.flush() self._last_progress_len = len(text) self._active_progress = True def finalize_progress_line(self, final_text: Optional[str] = None): with self._progress_lock: if final_text is not None: width = self._get_terminal_width() max_len = max(1, width - 1) if len(final_text) > max_len: final_text = final_text[:max_len] # Write final text, clear remainder, then newline sys.stdout.write("\r" + final_text) extra_spaces = max(0, max(self._last_progress_len - len(final_text), 0)) if extra_spaces: sys.stdout.write(" " * extra_spaces) sys.stdout.write("\r" + final_text) sys.stdout.write("\n") sys.stdout.flush() else: if self._active_progress: sys.stdout.write("\n") sys.stdout.flush() self._last_progress_len = 0 self._active_progress = False def log(self, msg: str): # Print a normal log line, ensuring it doesn't collide with progress line with self._progress_lock: if self._active_progress: self._clear_progress_line_locked() print(msg, flush=True) # --------------------------- SHA256 cache utils --------------------------- def load_sha_cache(self) -> Dict[str, str]: """Load SHA256 cache from file""" if os.path.exists(self.sha_cache_file): try: with open(self.sha_cache_file, 'r') as f: return json.load(f) except: return {} return {} def save_sha_cache(self): """Save SHA256 cache to file""" try: with open(self.sha_cache_file, 'w') as f: json.dump(self.sha_cache, f, indent=2) except Exception as e: self.log(f"Warning: Could not save SHA cache: {e}") def load_verified_cache(self) -> Dict[str, Dict]: """Load verified files cache from file""" if os.path.exists(self.verified_cache_file): try: with open(self.verified_cache_file, 'r') as f: return json.load(f) except: return {} return {} def save_verified_cache(self): """Save verified files cache to file""" try: with open(self.verified_cache_file, 'w') as f: json.dump(self.verified_cache, f, indent=2) print(f"[DEBUG] Verified cache saved to: {self.verified_cache_file}") print(f"[DEBUG] Cache contains {len(self.verified_cache)} entries") except Exception as e: self.log(f"Warning: Could not save verified cache: {e}") def is_file_verified(self, repo_id: str, filename: str, filepath: str, expected_sha: str) -> bool: """Check if file has already been verified and hasn't changed""" if not expected_sha: return False cache_key = f"{repo_id}/{filename}" if cache_key not in self.verified_cache: return False cached_info = self.verified_cache[cache_key] # Check if file exists and has same size and modification time if not os.path.exists(filepath): return False current_size = os.path.getsize(filepath) current_mtime = os.path.getmtime(filepath) return (cached_info.get('sha256') == expected_sha and cached_info.get('size') == current_size and abs(cached_info.get('mtime', 0) - current_mtime) < 1.0) # Allow 1 second tolerance def mark_file_verified(self, repo_id: str, filename: str, filepath: str, sha256: str): """Mark file as verified in cache""" cache_key = f"{repo_id}/{filename}" print(f"[DEBUG] Marking file as verified: {filename}") if os.path.exists(filepath): file_size = os.path.getsize(filepath) file_mtime = os.path.getmtime(filepath) self.verified_cache[cache_key] = { 'sha256': sha256, 'size': file_size, 'mtime': file_mtime, 'verified_at': time.time() } print(f"[DEBUG] Added to cache: {cache_key}") self.save_verified_cache() else: print(f"[DEBUG] File does not exist, cannot mark as verified: {filepath}") def get_file_sha256(self, repo_id: str, filename: str) -> Optional[str]: """Get SHA256 hash for a file from Hugging Face""" cache_key = f"{repo_id}/{filename}" # Check cache first if cache_key in self.sha_cache: return self.sha_cache[cache_key] try: # Method 1: Try to get from file metadata url = hf_hub_url(repo_id, filename) metadata = get_hf_file_metadata(url) # The etag contains the SHA256 hash if hasattr(metadata, 'etag'): # Remove quotes and 'W/' prefix if present sha256 = metadata.etag.replace('"', '').replace('W/', '') if len(sha256) == 64: # Valid SHA256 length self.sha_cache[cache_key] = sha256 self.save_sha_cache() return sha256 # Method 2: Try API approach api = HfApi() model_info = api.model_info(repo_id, files_metadata=True) for file_info in model_info.siblings: if file_info.rfilename == filename: if hasattr(file_info, 'lfs') and file_info.lfs: sha256 = file_info.lfs.get('sha256') if sha256: self.sha_cache[cache_key] = sha256 self.save_sha_cache() return sha256 return None except Exception as e: self.log(f"Warning: Could not get SHA256 for {filename}: {e}") return None # ----------------------- Formatting helpers (unchanged) ------------------- def format_bytes(self, bytes_val): """Format bytes to human readable format""" if bytes_val < 0: return "0 B" for unit in ['B', 'KB', 'MB', 'GB']: if bytes_val < 1024.0: return f"{bytes_val:.1f} {unit}" bytes_val /= 1024.0 return f"{bytes_val:.1f} TB" def format_time(self, seconds): """Format seconds to human readable format""" if seconds < 0: return "0s" seconds = int(seconds) # Convert to integer to avoid floating point issues if seconds < 60: return f"{seconds}s" elif seconds < 3600: mins = seconds // 60 secs = seconds % 60 if secs == 0: return f"{mins}m" else: return f"{mins}m {secs}s" else: hours = seconds // 3600 mins = (seconds % 3600) // 60 if mins == 0: return f"{hours}h" else: return f"{hours}h {mins}m" # ------------------------- Progress line rendering ------------------------ def print_progress(self, current, total, start_time, filename, speed_bytes_per_sec=None): """Render a single-line, non-wrapping progress line for the active file""" if total <= 0: return elapsed = max(0.001, time.time() - start_time) percent = min(100.0, (current / total * 100)) # Calculate speed if not provided if speed_bytes_per_sec is None: speed_bytes_per_sec = current / elapsed if elapsed > 0 else 0.0 # Calculate ETA if speed_bytes_per_sec > 0 and total > current: eta = (total - current) / speed_bytes_per_sec eta_str = self.format_time(eta) else: eta_str = "Complete" if current >= total else "Unknown" # Build the line # We will ensure the final line does not wrap by truncating to terminal width. # Progress bar target length (will be truncated by terminal width logic if needed) bar_length = 40 filled = int(percent * bar_length / 100) bar = "=" * max(0, min(filled, bar_length)) + "-" * max(0, bar_length - filled) speed_str = self.format_bytes(speed_bytes_per_sec) + "/s" if speed_bytes_per_sec else "0 B/s" line = ( f"{filename}: [{bar}] {percent:.1f}% " f"({self.format_bytes(current)}/{self.format_bytes(total)}) " f"{speed_str} ETA: {eta_str}" ) # Show single-line progress (truncated to fit, padded/cleared safely) self.show_progress_line(line) # ------------------------------ Networking -------------------------------- def get_file_url(self, repo_id: str, filename: str) -> str: """Get direct download URL for HuggingFace file""" return f"https://huggingface.co/{repo_id}/resolve/main/{filename}" def get_file_size(self, url: str) -> Optional[int]: """Get file size from remote server""" try: response = self.session.head(url, timeout=30, allow_redirects=True) if response.status_code == 200: content_length = response.headers.get('content-length') if content_length: return int(content_length) # Fallback to range request headers = {'Range': 'bytes=0-0'} response = self.session.get(url, headers=headers, timeout=30, stream=True) if response.status_code == 206: content_range = response.headers.get('content-range') if content_range and '/' in content_range: total_size = content_range.split('/')[-1] if total_size.isdigit(): return int(total_size) # Final fallback: check if file is likely small and compressed # Look for gzip encoding or chunked transfer which indicates small files if (response.headers.get('content-encoding') == 'gzip' or response.headers.get('transfer-encoding') == 'chunked'): # For compressed/chunked files, we'll return a special marker # to indicate we need to download the full file to get size return -1 return None except Exception as e: self.log(f"Warning: Could not get file size: {e}") return None # ----------------------- Chunk I/O and verification ----------------------- def verify_file_sha256(self, filepath: str, expected_sha: str, filename: str = "") -> bool: """Verify file SHA256 hash with single-line progress""" if not expected_sha: self.log(f"[WARNING] No SHA256 available for verification") return True # Can't verify, assume OK display_name = filename or os.path.basename(filepath) or filepath self.log(f"[VERIFYING] Computing SHA256 for {display_name}...") try: sha256_hash = hashlib.sha256() file_size = os.path.getsize(filepath) bytes_read = 0 start_time = time.time() last_update = 0.0 # Use 8MB chunks for better performance chunk_size = 8 * 1024 * 1024 with open(filepath, "rb") as f: while True: chunk = f.read(chunk_size) if not chunk: break sha256_hash.update(chunk) bytes_read += len(chunk) # Update progress once per ~0.2s to stay responsive now = time.time() if now - last_update >= 0.2: percent = (bytes_read / file_size) * 100 if file_size > 0 else 0.0 speed = bytes_read / max(0.001, (now - start_time)) line = ( f"[VERIFYING] {display_name}: {percent:.1f}% " f"({self.format_bytes(bytes_read)}/{self.format_bytes(file_size)}) " f"{self.format_bytes(speed)}/s" ) self.show_progress_line(line) last_update = now # Finalize computed_sha = sha256_hash.hexdigest() if computed_sha == expected_sha: self.finalize_progress_line(f"[VERIFIED] SHA256 match: {computed_sha[:16]}...") return True else: self.finalize_progress_line(f"[ERROR] SHA256 mismatch!") self.log(f" Expected: {expected_sha}") self.log(f" Got: {computed_sha}") return False except Exception as e: self.finalize_progress_line() self.log(f"[ERROR] Failed to verify SHA256: {e}") return False def download_chunk(self, url: str, start: int, end: int, filepath: str, chunk_id: int, progress_callback=None) -> bool: """Download a specific chunk with resume support""" chunk_file = os.path.normpath(f"{filepath}.part{chunk_id}") chunk_size_expected = end - start + 1 # Check if chunk already complete if os.path.exists(chunk_file): existing_size = os.path.getsize(chunk_file) if existing_size == chunk_size_expected: if progress_callback: progress_callback(chunk_id, chunk_size_expected) return True elif existing_size > chunk_size_expected: # Corrupted chunk, delete and restart os.remove(chunk_file) resume_pos = 0 else: # Valid partial chunk, resume resume_pos = existing_size else: resume_pos = 0 max_retries = self.config["max_retries"] for attempt in range(max_retries): try: # Download from resume position actual_start = start + resume_pos headers = {'Range': f'bytes={actual_start}-{end}'} response = self.session.get(url, headers=headers, timeout=self.config["timeout"], stream=True) if response.status_code not in [200, 206]: raise Exception(f"Bad status code: {response.status_code}") # Download chunk mode = 'ab' if resume_pos > 0 else 'wb' downloaded = resume_pos with open(chunk_file, mode) as f: for data in response.iter_content(chunk_size=self.config["chunk_size"]): if data: f.write(data) downloaded += len(data) if progress_callback: progress_callback(chunk_id, downloaded) # Verify chunk is complete final_size = os.path.getsize(chunk_file) if final_size == chunk_size_expected: if progress_callback: progress_callback(chunk_id, chunk_size_expected) return True elif final_size < chunk_size_expected: # Incomplete, will retry resume_pos = final_size raise Exception(f"Chunk incomplete: {final_size}/{chunk_size_expected}") else: # Too large, corrupted os.remove(chunk_file) resume_pos = 0 raise Exception(f"Chunk too large: {final_size}/{chunk_size_expected}") except Exception as e: if attempt < max_retries - 1: delay = min(self.config["retry_delay"] * (2 ** attempt), self.config["max_retry_delay"]) time.sleep(delay) else: self.log(f"Chunk {chunk_id} failed after {max_retries} attempts: {e}") return False return False # ------------------------------ Merge chunks ------------------------------ def merge_chunks(self, filepath: str, num_chunks: int) -> bool: """Merge downloaded chunks into final file with optimized I/O""" temp_file = os.path.normpath(f"{filepath}.tmp") try: # Check all chunks exist and calculate total size total_size = 0 chunk_files = [] for i in range(num_chunks): chunk_file = os.path.normpath(f"{filepath}.part{i}") if not os.path.exists(chunk_file): self.log(f"Error: Missing chunk {i}") return False size = os.path.getsize(chunk_file) chunk_files.append((chunk_file, size)) total_size += size # For very large files (>1GB), use optimized OS-level copy if total_size > 1024 * 1024 * 1024: # 1GB return self.merge_chunks_optimized(filepath, chunk_files, total_size, temp_file) # For smaller files, use buffered I/O buffer_size = 64 * 1024 * 1024 # 64MB buffer merged_size = 0 start_time = time.time() last_update = 0.0 with open(temp_file, 'wb') as outfile: for chunk_file, chunk_size in chunk_files: with open(chunk_file, 'rb') as infile: bytes_copied = 0 while bytes_copied < chunk_size: to_read = min(buffer_size, chunk_size - bytes_copied) data = infile.read(to_read) if not data: break outfile.write(data) bytes_copied += len(data) merged_size += len(data) # Update progress once per ~0.5s now = time.time() if now - last_update >= 0.5: percent = (merged_size / total_size) * 100 if total_size > 0 else 0.0 speed = merged_size / max(0.001, (now - start_time)) line = ( f"[MERGING] {percent:.1f}% " f"({self.format_bytes(merged_size)}/{self.format_bytes(total_size)}) " f"Speed: {self.format_bytes(speed)}/s" ) self.show_progress_line(line) last_update = now # Final progress line elapsed = max(0.001, time.time() - start_time) avg_speed = total_size / elapsed self.finalize_progress_line( f"[MERGING] 100.0% ({self.format_bytes(total_size)}/{self.format_bytes(total_size)}) " f"Completed in {self.format_time(elapsed)} - Avg: {self.format_bytes(avg_speed)}/s" ) # Replace target file if os.path.exists(filepath): os.remove(filepath) os.rename(temp_file, filepath) # Clean up chunks after successful merge for chunk_file, _ in chunk_files: try: os.remove(chunk_file) except: pass return True except Exception as e: self.finalize_progress_line() self.log(f"Error merging: {e}") if os.path.exists(temp_file): try: os.remove(temp_file) except: pass return False def merge_chunks_optimized(self, filepath: str, chunk_files: List[Tuple[str, int]], total_size: int, temp_file: str) -> bool: """Optimized merge for very large files using OS-level operations""" try: self.log(f"[MERGING] Using optimized merge for large file ({self.format_bytes(total_size)})") start_time = time.time() merged_size = 0 last_update = 0.0 # Use buffered concat for portability with open(temp_file, 'wb') as outfile: for i, (chunk_file, chunk_size) in enumerate(chunk_files): # Update line per chunk and overall bytes to avoid extra prints now = time.time() if now - last_update >= 0.5: line = f"[MERGING] Chunk {i+1}/{len(chunk_files)} " \ f"({self.format_bytes(merged_size)}/{self.format_bytes(total_size)})" self.show_progress_line(line) last_update = now with open(chunk_file, 'rb') as infile: shutil.copyfileobj(infile, outfile, length=128 * 1024 * 1024) # 128MB buffer merged_size += chunk_size elapsed = max(0.001, time.time() - start_time) avg_speed = total_size / elapsed self.finalize_progress_line( f"[MERGING] 100.0% - Completed in {self.format_time(elapsed)} - Avg: {self.format_bytes(avg_speed)}/s" ) # Replace target file if os.path.exists(filepath): os.remove(filepath) os.rename(temp_file, filepath) # Clean up chunks for chunk_file, _ in chunk_files: try: os.remove(chunk_file) except: pass return True except Exception as e: self.finalize_progress_line() self.log(f"Error in optimized merge: {e}") if os.path.exists(temp_file): try: os.remove(temp_file) except: pass return False # ------------------------------ Download API ------------------------------ def download_file(self, repo_id: str, filename: str, local_dir: str) -> bool: """Main download function with SHA256 verification""" url = self.get_file_url(repo_id, filename) filepath = os.path.normpath(os.path.join(local_dir, filename)) # Ensure the directory for the file exists (handle subdirectories) file_dir = os.path.dirname(filepath) os.makedirs(file_dir, exist_ok=True) # Get expected SHA256 expected_sha = self.get_file_sha256(repo_id, filename) if expected_sha: self.log(f"[INFO] Expected SHA256: {expected_sha[:16]}...") # Get file size file_size = self.get_file_size(url) # Check if already complete if os.path.exists(filepath): actual_size = os.path.getsize(filepath) if file_size: if actual_size == file_size: # Check if already verified to skip SHA256 computation if expected_sha and self.is_file_verified(repo_id, filename, filepath, expected_sha): self.log(f"[SKIP] {filename} already complete and verified (cached) ({self.format_bytes(file_size)})") return True # Verify SHA256 if available and not cached elif expected_sha: if self.verify_file_sha256(filepath, expected_sha, filename): self.mark_file_verified(repo_id, filename, filepath, expected_sha) self.log(f"[SKIP] {filename} already complete and verified ({self.format_bytes(file_size)})") return True else: self.log(f"[WARNING] {filename} failed verification, re-downloading") os.remove(filepath) else: self.log(f"[SKIP] {filename} already complete ({self.format_bytes(file_size)})") return True elif actual_size > file_size: self.log(f"[WARNING] {filename} corrupted, re-downloading") os.remove(filepath) else: # Can't verify size, try SHA256 verification if actual_size > 1024 and expected_sha: # > 1KB # Check if already verified to skip SHA256 computation if self.is_file_verified(repo_id, filename, filepath, expected_sha): self.log(f"[SKIP] {filename} exists and verified (cached) ({self.format_bytes(actual_size)})") return True elif self.verify_file_sha256(filepath, expected_sha, filename): self.mark_file_verified(repo_id, filename, filepath, expected_sha) self.log(f"[SKIP] {filename} exists and verified ({self.format_bytes(actual_size)})") return True else: self.log(f"[WARNING] {filename} failed verification, re-downloading") os.remove(filepath) elif actual_size > 1024: # No SHA to verify self.log(f"[SKIP] {filename} exists ({self.format_bytes(actual_size)})") return True if file_size is None: self.log(f"[ERROR] Cannot get size for {filename}") return False elif file_size == -1: # Handle compressed/chunked files - download without size info self.log(f"[INFO] {filename} is compressed, downloading without size info") return self.download_unknown_size(url, filepath, filename, expected_sha, repo_id) # Use parallel download for files > 10MB if file_size > 10 * 1024 * 1024: success = self.download_parallel(url, filepath, filename, file_size) else: success = self.download_single(url, filepath, filename, file_size) # Verify SHA256 after successful download if success and expected_sha: if not self.verify_file_sha256(filepath, expected_sha, filename): self.log(f"[ERROR] {filename} downloaded but failed SHA256 verification") try: os.remove(filepath) except: pass return False else: # Mark file as verified after successful verification self.mark_file_verified(repo_id, filename, filepath, expected_sha) return success def download_unknown_size(self, url: str, filepath: str, filename: str, expected_sha: str, repo_id: str = "") -> bool: """Download file when size cannot be determined (compressed/chunked files)""" max_retries = self.config["max_retries"] for attempt in range(max_retries): try: self.log(f"[DOWNLOADING] {filename} (unknown size)") response = self.session.get(url, timeout=self.config["timeout"], stream=True) response.raise_for_status() # Download the file downloaded = 0 start_time = time.time() last_update = 0.0 with open(filepath, 'wb') as f: for chunk in response.iter_content(chunk_size=self.config["chunk_size"]): if chunk: f.write(chunk) downloaded += len(chunk) # Show progress without total size now = time.time() if now - last_update >= 0.5: elapsed = max(0.001, now - start_time) speed = downloaded / elapsed line = (f"[DOWNLOADING] {filename}: {self.format_bytes(downloaded)} " f"@ {self.format_bytes(speed)}/s") self.show_progress_line(line) last_update = now # Finalize download final_size = os.path.getsize(filepath) elapsed = max(0.001, time.time() - start_time) avg_speed = final_size / elapsed self.finalize_progress_line( f"[OK] {filename} completed ({self.format_bytes(final_size)}) " f"in {self.format_time(elapsed)} - Avg: {self.format_bytes(avg_speed)}/s" ) # Verify SHA256 if available if expected_sha: if not self.verify_file_sha256(filepath, expected_sha, filename): self.log(f"[ERROR] {filename} downloaded but failed SHA256 verification") try: os.remove(filepath) except: pass return False else: # Mark file as verified after successful verification if repo_id: self.mark_file_verified(repo_id, filename, filepath, expected_sha) return True except Exception as e: self.finalize_progress_line() self.log(f"[ERROR] Attempt {attempt + 1}: {e}") if attempt < max_retries - 1: delay = min(self.config["retry_delay"] * (2 ** attempt), self.config["max_retry_delay"]) time.sleep(delay) else: return False return False def download_parallel(self, url: str, filepath: str, filename: str, file_size: int) -> bool: """Download using 16 parallel connections""" num_chunks = self.config["num_connections"] # Calculate chunk boundaries base_chunk_size = file_size // num_chunks chunks = [] for i in range(num_chunks): start = i * base_chunk_size if i == num_chunks - 1: # Last chunk gets all remaining bytes end = file_size - 1 else: end = (i + 1) * base_chunk_size - 1 chunks.append((i, start, end)) self.log(f"[DOWNLOADING] {filename} ({self.format_bytes(file_size)}) using {num_chunks} connections") # Progress tracking chunk_progress = {} progress_lock = threading.Lock() def update_progress(chunk_id: int, bytes_downloaded: int): with progress_lock: chunk_progress[chunk_id] = bytes_downloaded # Check existing progress for chunk_id, start, end in chunks: chunk_file = os.path.normpath(f"{filepath}.part{chunk_id}") if os.path.exists(chunk_file): size = os.path.getsize(chunk_file) chunk_progress[chunk_id] = size else: chunk_progress[chunk_id] = 0 initial_bytes = sum(chunk_progress.values()) if initial_bytes > 0: self.log(f"[RESUMING] Already downloaded: {self.format_bytes(initial_bytes)}") # Download chunks start_time = time.time() failed_chunks = [] with concurrent.futures.ThreadPoolExecutor(max_workers=num_chunks) as executor: # Submit all chunks futures = {} for chunk_id, start, end in chunks: # Only download if not complete chunk_file = os.path.normpath(f"{filepath}.part{chunk_id}") expected_size = end - start + 1 if os.path.exists(chunk_file) and os.path.getsize(chunk_file) == expected_size: continue # Skip completed chunks future = executor.submit( self.download_chunk, url, start, end, filepath, chunk_id, update_progress ) futures[future] = chunk_id # If all chunks already complete if not futures: self.log("[MERGING] All chunks already complete") if self.merge_chunks(filepath, num_chunks): # Verify final file if os.path.getsize(filepath) == file_size: self.log(f"[OK] {filename} completed") return True else: self.log(f"[ERROR] Size mismatch after merge") try: os.remove(filepath) except: pass return False return False # Monitor progress last_update = 0.0 last_bytes = initial_bytes while futures: # Wait for any future to complete done, pending = concurrent.futures.wait( list(futures.keys()), timeout=0.5, return_when=concurrent.futures.FIRST_COMPLETED ) # Process completed futures for future in done: chunk_id = futures.pop(future) try: success = future.result() if not success: failed_chunks.append(chunk_id) except Exception as e: self.log(f"Chunk {chunk_id} exception: {e}") failed_chunks.append(chunk_id) # Update progress (once per ~1s) current_time = time.time() if current_time - last_update >= 1.0 or not futures: with progress_lock: current_bytes = sum(chunk_progress.values()) # Calculate recent speed time_delta = current_time - last_update if last_update > 0 else (current_time - start_time) bytes_delta = current_bytes - last_bytes if last_update > 0 else (current_bytes - initial_bytes) speed = bytes_delta / max(0.001, time_delta) self.print_progress(current_bytes, file_size, start_time, filename, speed) last_update = current_time last_bytes = current_bytes # Remove the active progress line so the next messages appear cleanly self.clear_progress_line() # Check results if failed_chunks: self.log(f"[ERROR] Failed chunks: {failed_chunks}") return False # Verify all chunks are complete before merging for chunk_id, start, end in chunks: chunk_file = os.path.normpath(f"{filepath}.part{chunk_id}") expected_size = end - start + 1 if not os.path.exists(chunk_file): self.log(f"[ERROR] Missing chunk {chunk_id}") return False actual_size = os.path.getsize(chunk_file) if actual_size != expected_size: self.log(f"[ERROR] Chunk {chunk_id} incomplete: {actual_size}/{expected_size}") return False # Merge chunks self.log(f"[MERGING] Merging {num_chunks} chunks...") if self.merge_chunks(filepath, num_chunks): # Verify final file size final_size = os.path.getsize(filepath) if final_size == file_size: elapsed = max(0.001, time.time() - start_time) avg_speed = (file_size - initial_bytes) / elapsed if elapsed > 0 else 0 self.log(f"[OK] {filename} completed in {self.format_time(elapsed)} " f"- Average: {self.format_bytes(avg_speed)}/s") return True else: self.log(f"[ERROR] Final size mismatch: {final_size} != {file_size}") try: os.remove(filepath) except: pass return False else: self.log(f"[ERROR] Merge failed") return False def download_single(self, url: str, filepath: str, filename: str, file_size: int) -> bool: """Single connection download for small files""" max_retries = self.config["max_retries"] for attempt in range(max_retries): try: # Check existing file resume_pos = 0 if os.path.exists(filepath): resume_pos = os.path.getsize(filepath) if resume_pos >= file_size: self.log(f"[OK] {filename} already complete") return True # Download headers = {'Range': f'bytes={resume_pos}-'} if resume_pos > 0 else {} response = self.session.get(url, headers=headers, timeout=self.config["timeout"], stream=True) if resume_pos > 0 and response.status_code != 206: self.log(f"[WARNING] Resume not supported, restarting") resume_pos = 0 response = self.session.get(url, timeout=self.config["timeout"], stream=True) response.raise_for_status() # Write file mode = 'ab' if resume_pos > 0 else 'wb' downloaded = 0 start_time = time.time() last_update = 0.0 if resume_pos > 0: self.log(f"[RESUMING] {filename} from {self.format_bytes(resume_pos)}") else: self.log(f"[DOWNLOADING] {filename} ({self.format_bytes(file_size)})") with open(filepath, mode) as f: for chunk in response.iter_content(chunk_size=self.config["chunk_size"]): if chunk: f.write(chunk) downloaded += len(chunk) # Progress update once per ~0.5s now = time.time() if now - last_update >= 0.5: total = resume_pos + downloaded self.print_progress(total, file_size, start_time, filename) last_update = now # Final progress update total = resume_pos + downloaded self.print_progress(total, file_size, start_time, filename) # Move to next line cleanly self.clear_progress_line() # Verify size if os.path.getsize(filepath) == file_size: self.log(f"[OK] {filename} completed") return True else: self.log(f"[ERROR] Size mismatch") continue except Exception as e: self.finalize_progress_line() self.log(f"[ERROR] Attempt {attempt + 1}: {e}") if attempt < max_retries - 1: delay = min(self.config["retry_delay"] * (2 ** attempt), self.config["max_retry_delay"]) time.sleep(delay) else: return False return False # ------------------------------- Repo scanning ------------------------------- def scan_repo_files(repo_id: str, specific_files: List[str] = None) -> List[str]: """Get list of files to download""" try: if specific_files: print(f"Using predefined file list") return specific_files api = HfApi() all_files = list_repo_files(repo_id=repo_id, repo_type="model") return all_files except Exception as e: print(f"Error scanning repository: {e}") return specific_files if specific_files else [] def parse_model_file_specs(file_specs) -> List[Tuple[str, str]]: """ Normalize model config file specs to a list of (remote_path, local_filename). Supported formats: - "filename.ext" (remote == local) - {"remote": "path/in/repo.ext", "local": "local_name.ext"} - ("remote_path", "local_name") """ items: List[Tuple[str, str]] = [] for spec in file_specs or []: if isinstance(spec, str): items.append((spec, spec)) continue if isinstance(spec, dict): remote = spec.get("remote") local = spec.get("local") if not remote: raise ValueError(f"Invalid file spec (missing 'remote'): {spec}") if not local: local = os.path.basename(remote) items.append((remote, local)) continue if isinstance(spec, (list, tuple)) and len(spec) == 2: remote, local = spec items.append((str(remote), str(local))) continue raise ValueError(f"Unsupported file spec format: {spec}") seen = set() deduped: List[Tuple[str, str]] = [] for remote, local in items: key = (remote, local) if key in seen: continue seen.add(key) deduped.append((remote, local)) return deduped def download_file_with_rename(downloader: RobustDownloader, repo_id: str, remote_filename: str, local_dir: str, local_filename: str) -> bool: """Download a file from HuggingFace repo but save it with a different local name""" url = downloader.get_file_url(repo_id, remote_filename) filepath = os.path.normpath(os.path.join(local_dir, local_filename)) # Ensure the directory for the file exists (handle subdirectories) file_dir = os.path.dirname(filepath) os.makedirs(file_dir, exist_ok=True) # Get expected SHA256 - using the original remote filename expected_sha = downloader.get_file_sha256(repo_id, remote_filename) if expected_sha: downloader.log(f"[INFO] Expected SHA256: {expected_sha[:16]}...") # Get file size file_size = downloader.get_file_size(url) # Check if already complete if os.path.exists(filepath): actual_size = os.path.getsize(filepath) if file_size: if actual_size == file_size: # Check if already verified to skip SHA256 computation if expected_sha and downloader.is_file_verified(repo_id, remote_filename, filepath, expected_sha): downloader.log(f"[SKIP] {local_filename} already complete and verified (cached) ({downloader.format_bytes(file_size)})") return True # Verify SHA256 if available and not cached elif expected_sha: if downloader.verify_file_sha256(filepath, expected_sha, local_filename): downloader.mark_file_verified(repo_id, remote_filename, filepath, expected_sha) downloader.log(f"[SKIP] {local_filename} already complete and verified ({downloader.format_bytes(file_size)})") return True else: downloader.log(f"[WARNING] {local_filename} failed verification, re-downloading") os.remove(filepath) else: downloader.log(f"[SKIP] {local_filename} already complete ({downloader.format_bytes(file_size)})") return True elif actual_size > file_size: downloader.log(f"[WARNING] {local_filename} corrupted, re-downloading") os.remove(filepath) else: # Can't verify size, try SHA256 verification if actual_size > 1024 and expected_sha: # > 1KB # Check if already verified to skip SHA256 computation if downloader.is_file_verified(repo_id, remote_filename, filepath, expected_sha): downloader.log(f"[SKIP] {local_filename} exists and verified (cached) ({downloader.format_bytes(actual_size)})") return True elif downloader.verify_file_sha256(filepath, expected_sha, local_filename): downloader.mark_file_verified(repo_id, remote_filename, filepath, expected_sha) downloader.log(f"[SKIP] {local_filename} exists and verified ({downloader.format_bytes(actual_size)})") return True else: downloader.log(f"[WARNING] {local_filename} failed verification, re-downloading") os.remove(filepath) elif actual_size > 1024: # No SHA to verify downloader.log(f"[SKIP] {local_filename} exists ({downloader.format_bytes(actual_size)})") return True if file_size is None: downloader.log(f"[ERROR] Cannot get size for {remote_filename}") return False elif file_size == -1: # Handle compressed/chunked files - download without size info downloader.log(f"[INFO] {local_filename} is compressed, downloading without size info") return downloader.download_unknown_size(url, filepath, local_filename, expected_sha, repo_id) # Use parallel download for files > 10MB if file_size > 10 * 1024 * 1024: success = downloader.download_parallel(url, filepath, local_filename, file_size) else: success = downloader.download_single(url, filepath, local_filename, file_size) # Verify SHA256 after successful download if success and expected_sha: if not downloader.verify_file_sha256(filepath, expected_sha, local_filename): downloader.log(f"[ERROR] {local_filename} downloaded but failed SHA256 verification") try: os.remove(filepath) except: pass return False else: # Mark file as verified after successful verification downloader.mark_file_verified(repo_id, remote_filename, filepath, expected_sha) return success def _get_terminal_width(fallback: int = 140) -> int: try: return shutil.get_terminal_size(fallback=(fallback, 20)).columns except Exception: return fallback def _shorten_end(text: str, max_len: int) -> str: if max_len <= 0: return "" if len(text) <= max_len: return text if max_len <= 3: return text[:max_len] return text[: max_len - 3] + "..." def _shorten_middle(text: str, max_len: int) -> str: if max_len <= 0: return "" if len(text) <= max_len: return text if max_len <= 3: return text[:max_len] left = (max_len - 3) // 2 right = (max_len - 3) - left return text[:left] + "..." + text[-right:] def render_available_models_table() -> str: entries = [] for i, model_id in enumerate(MODEL_MENU, 1): cfg = MODEL_CONFIGS.get(model_id, {}) files = [local for _, local in parse_model_file_specs(cfg.get("files", []))] entries.append( { "choice": str(i), "model_id": model_id, "name": cfg.get("name", model_id), "default_dir": cfg.get("default_dir", ""), "files": files, } ) headers = ["Choice", "Training Bundle", "Downloaded Models"] choice_w = max(len(headers[0]), max(len(e["choice"]) for e in entries)) max_bundle_len = max(len(e["name"]) for e in entries) term_w = _get_terminal_width() min_bundle_w = max(len(headers[1]), 22) min_models_w = max(len(headers[2]), 30) available = term_w - choice_w - 10 if available < (min_bundle_w + min_models_w): # Ensure we still render a table even in very narrow terminals min_bundle_w = max(14, min_bundle_w) min_models_w = max(18, min_models_w) bundle_w = min(max_bundle_len, max(min_bundle_w, int(available * 0.35))) models_w = max(min_models_w, available - bundle_w) if bundle_w + models_w > available: models_w = max(min_models_w, available - bundle_w) if models_w < min_models_w: models_w = min_models_w bundle_w = max(min_bundle_w, available - models_w) if bundle_w < min_bundle_w: bundle_w = min_bundle_w models_w = max(min_models_w, available - bundle_w) def wrap_cell(text: str, width: int) -> List[str]: if width <= 0: return [""] if not text: return [""] return textwrap.wrap(text, width=width, break_long_words=True, break_on_hyphens=False) or [""] def border_line() -> str: return "+" + "-" * (choice_w + 2) + "+" + "-" * (bundle_w + 2) + "+" + "-" * (models_w + 2) + "+" def format_row(choice: str, bundle: str, models: str) -> List[str]: bundle_lines = wrap_cell(bundle, bundle_w) model_lines = wrap_cell(models, models_w) height = max(len(bundle_lines), len(model_lines), 1) lines = [] for i in range(height): ch = choice if i == 0 else "" b = bundle_lines[i] if i < len(bundle_lines) else "" m = model_lines[i] if i < len(model_lines) else "" lines.append(f"| {ch:<{choice_w}} | {b:<{bundle_w}} | {m:<{models_w}} |") return lines lines = [border_line()] lines.extend(format_row(headers[0], headers[1], headers[2])) lines.append(border_line()) for e in entries: models_text = ", ".join(e["files"]) if e["files"] else "-" lines.extend(format_row(e["choice"], e["name"], models_text)) lines.append(border_line()) return "\n".join(lines).rstrip() def get_model_choice(): """Get model choice from user""" print("\nAvailable models:") print(render_available_models_table()) print() choice_map = {str(i): model_id for i, model_id in enumerate(MODEL_MENU, 1)} max_choice = len(MODEL_MENU) while True: try: choice = input(f"Please select model (1-{max_choice}): ").strip() if choice in choice_map: return choice_map[choice] print(f"Invalid choice. Please enter a number from 1 to {max_choice}.") except KeyboardInterrupt: print("\nDownload cancelled.") sys.exit(0) def normalize_path(path_str: str) -> str: """Normalize path to work on both Windows and Linux""" if not path_str: return path_str # Convert to Path object and resolve path = Path(path_str).expanduser() # If it's not absolute, make it relative to current directory if not path.is_absolute(): path = Path.cwd() / path # Normalize and return as string return str(path.resolve()) def download_models(model_id: Optional[str] = None, download_dir: Optional[str] = None): """Main download function""" # Get model choice if not provided if not model_id: model_id = get_model_choice() # Validate model_id if model_id not in MODEL_CONFIGS: print(f"Error: Invalid model ID '{model_id}'. Available options: {list(MODEL_CONFIGS.keys())}") return model_config = MODEL_CONFIGS[model_id] # Determine target directory if download_dir: target_dir = normalize_path(download_dir) else: target_dir = normalize_path(model_config["default_dir"]) repo_id = model_config["repo_id"] download_items = parse_model_file_specs(model_config["files"]) print(f"\nModel: {model_config['name']}") print(f"Description: {model_config['description']}") downloader = RobustDownloader(DOWNLOAD_CONFIG) os.makedirs(target_dir, exist_ok=True) print(f"\nRepository: {repo_id}") print(f"Target: {os.path.abspath(target_dir)}") scan_repo_files(repo_id, [remote for remote, _ in download_items]) if not download_items: print("No files found") return # Check existing files with SHA256 verification to_download = [] skipped = [] for remote_filename, local_filename in download_items: filepath = os.path.normpath(os.path.join(target_dir, local_filename)) if os.path.exists(filepath) and os.path.getsize(filepath) > 1024: # Try to verify with SHA256 expected_sha = downloader.get_file_sha256(repo_id, remote_filename) if expected_sha: # Check if already verified to skip SHA256 computation if downloader.is_file_verified(repo_id, remote_filename, filepath, expected_sha): print(f"[SKIP] {local_filename} already verified (cached)") skipped.append(local_filename) elif downloader.verify_file_sha256(filepath, expected_sha, local_filename): # Mark as verified after successful verification downloader.mark_file_verified(repo_id, remote_filename, filepath, expected_sha) skipped.append(local_filename) else: print(f"[RE-DOWNLOAD] {local_filename} failed verification") to_download.append((remote_filename, local_filename)) else: # No SHA available, skip if file exists skipped.append(local_filename) else: to_download.append((remote_filename, local_filename)) print(f"\nTotal files: {len(download_items)}") if skipped: print(f"Already downloaded: {len(skipped)}") for f in skipped: print(f" [OK] {f}") if not to_download: print("All files already downloaded!") return print(f"To download: {len(to_download)}") # Download files successful = 0 failed = 0 for i, (remote_filename, local_filename) in enumerate(to_download, 1): print(f"\n[{i}/{len(to_download)}]") if download_file_with_rename(downloader, repo_id, remote_filename, target_dir, local_filename): successful += 1 else: failed += 1 print(f"[FAILED] {local_filename}") # Summary print(f"\n{'='*50}") print(f"Summary:") print(f" Already had: {len(skipped)}") print(f" Downloaded: {successful}") print(f" Failed: {failed}") print(f"\n{'='*50}") print(f"Downloads complete!") print(f"Location: {os.path.abspath(target_dir)}") if __name__ == "__main__": parser = argparse.ArgumentParser(description='Download training models for Qwen Image or Wan Text-to-Video') parser.add_argument( '--model', choices=[ 'qwen_image', 'qwen_image_2512', 'qwen_image_edit_plus', 'qwen_image_edit_2511', 'wan21_t2v', 'wan22_t2v', 'wan22_i2v', 'flux2_dev', 'flux_klein_9b', 'flux_klein_4b', 'z_image_base', 'z_image_turbo', ], help=( 'Model to download: ' 'qwen_image (Qwen Image Training), ' 'qwen_image_2512 (Qwen Image 2512), ' 'qwen_image_edit_plus (Qwen Image Edit Plus 2509), ' 'qwen_image_edit_2511 (Qwen Image Edit 2511), ' 'wan21_t2v (Wan 2.1 Text-to-Video), ' 'wan22_t2v (Wan 2.2 Text-to-Video), ' 'wan22_i2v (Wan 2.2 Image-to-Video), ' 'flux2_dev (FLUX 2 Dev), ' 'flux_klein_9b (FLUX Klein 9B), ' 'flux_klein_4b (FLUX Klein 4B), ' 'z_image_base (Z-Image Base), ' 'or z_image_turbo (Z-Image Turbo)' ), ) parser.add_argument('--dir', type=str, help='Download directory (supports both full and relative paths for Windows and Linux)') parser.add_argument('--list', action='store_true', help='Print available models table and exit') args = parser.parse_args() if args.list: print(render_available_models_table()) sys.exit(0) download_models(args.model, args.dir)