"""Helper utilities for Z-Image.""" import hashlib import json from pathlib import Path from typing import Optional, List, Tuple, Dict from loguru import logger import torch from config import BYTES_PER_GB def format_bytes(size: float) -> str: """ Format bytes to GB string. Args: size: Size in bytes Returns: Formatted string in GB """ n = size / BYTES_PER_GB return f"{n:.2f} GB" def print_memory_stats(stage: str) -> None: """ Print CUDA memory statistics. Args: stage: Description of current stage """ if not torch.cuda.is_available(): logger.warning("CUDA not available, skipping memory stats") return torch.cuda.synchronize() allocated = torch.cuda.max_memory_allocated() reserved = torch.cuda.max_memory_reserved() current_allocated = torch.cuda.memory_allocated() current_reserved = torch.cuda.memory_reserved() logger.info(f"[{stage}] Memory Stats:") logger.info(f" Current Allocated: {format_bytes(current_allocated)}") logger.info(f" Current Reserved: {format_bytes(current_reserved)}") logger.info(f" Peak Allocated: {format_bytes(allocated)}") logger.info(f" Peak Reserved: {format_bytes(reserved)}") def compute_file_md5(file_path: Path, chunk_size: int = 8192) -> str: """Compute MD5 hash of a file.""" md5_hash = hashlib.md5() with open(file_path, "rb") as f: while chunk := f.read(chunk_size): md5_hash.update(chunk) return md5_hash.hexdigest() def load_manifest(manifest_file: Path) -> Dict[str, Optional[str]]: """Load manifest file. Returns dict mapping file paths to MD5 hashes (or None).""" manifest = {} if not manifest_file.exists(): return manifest with open(manifest_file, "r", encoding="utf-8") as f: for line_num, line in enumerate(f, 1): line = line.strip() # Skip empty lines and comments if not line or line.startswith("#"): continue parts = line.split() if len(parts) == 1: # Only file path, no checksum file_path = parts[0] manifest[file_path] = None elif len(parts) == 2: # File path with checksum if len(parts[0]) == 32 and all(c in '0123456789abcdef' for c in parts[0].lower()): md5_hash, file_path = parts else: file_path, md5_hash = parts manifest[file_path] = md5_hash else: logger.warning(f"Invalid manifest format at line {line_num}: {line}") continue return manifest def verify_file_integrity( base_dir: Path, manifest: Dict[str, Optional[str]], verify_checksums: bool = True ) -> Tuple[bool, List[str], List[str]]: """ Verify file integrity using a manifest. Args: base_dir: Base directory for relative file paths manifest: Dictionary of relative paths to MD5 hashes (None if no hash provided) verify_checksums: If True, verify MD5 checksums when available; if False, only check existence Returns: Tuple of (all_valid: bool, missing_files: List[str], corrupted_files: List[str]) """ missing = [] corrupted = [] for rel_path, expected_md5 in manifest.items(): file_path = base_dir / rel_path if not file_path.exists(): missing.append(rel_path) continue # Only verify checksum if requested AND hash is available if verify_checksums and expected_md5 is not None: try: actual_md5 = compute_file_md5(file_path) if actual_md5 != expected_md5: corrupted.append(rel_path) logger.debug(f"Checksum mismatch for {rel_path}: expected {expected_md5}, got {actual_md5}") except Exception as e: logger.error(f"Failed to compute checksum for {rel_path}: {e}") corrupted.append(rel_path) all_valid = len(missing) == 0 and len(corrupted) == 0 return all_valid, missing, corrupted def ensure_model_weights( model_path: str, repo_id: str = "Tongyi-MAI/Z-Image-Turbo", verify: bool = False, manifest_name: Optional[str] = None ) -> Path: """ Ensure model weights exist and optionally verify integrity. Args: model_path: Path to model directory repo_id: HuggingFace repo ID for download verify: If True, verify MD5 checksums; if False, only check existence manifest_name: Manifest file name in src/config/manifests/ (auto-detect if None) Returns: Path to validated model directory """ from huggingface_hub import snapshot_download target_dir = Path(model_path) # Determine manifest path if manifest_name: # Explicitly specified manifest from config/manifests/ manifest_path = Path(__file__).parent.parent / "config" / "manifests" / manifest_name else: # Auto-detect model_name = target_dir.name.lower() # e.g., "Z-Image-Turbo" -> "z-image-turbo" config_manifest = Path(__file__).parent.parent / "config" / "manifests" / f"{model_name}.txt" if config_manifest.exists(): manifest_path = config_manifest else: # Fallback manifest_path = target_dir / "manifest.txt" manifest = load_manifest(manifest_path) if not manifest: logger.warning(f"Manifest file not found: {manifest_path}") logger.warning("Skipping file verification (assuming model exists)") if target_dir.exists(): logger.info(f"✓ Model directory exists: {target_dir}") return target_dir else: logger.warning(f"Model directory not found: {target_dir}") missing_files = ["entire model directory"] corrupted_files = [] else: # Count files with checksums files_with_checksums = sum(1 for v in manifest.values() if v is not None) if verify and files_with_checksums == 0: logger.info(f"Verify requested but no checksums in manifest, only checking existence") elif verify and files_with_checksums > 0: logger.info(f"Verifying {files_with_checksums} file(s) with MD5 checksums...") # Verify files all_valid, missing_files, corrupted_files = verify_file_integrity( target_dir, manifest, verify_checksums=verify ) if all_valid: if verify and files_with_checksums > 0: logger.success(f"✓ All files verified with MD5 checksums in {target_dir}") else: logger.info(f"✓ All {len(manifest)} required files exist in {target_dir}") return target_dir # Report missing and corrupted files if missing_files: logger.warning(f"Missing {len(missing_files)} file(s):") for f in missing_files[:10]: logger.warning(f" - {f}") if len(missing_files) > 10: logger.warning(f" ... and {len(missing_files) - 10} more") if corrupted_files: logger.error(f"Corrupted {len(corrupted_files)} file(s) (checksum mismatch):") for f in corrupted_files[:10]: logger.error(f" - {f}") if len(corrupted_files) > 10: logger.error(f" ... and {len(corrupted_files) - 10} more") # Download model weights logger.info(f"\nAttempting to download from {repo_id}...") try: target_dir.mkdir(parents=True, exist_ok=True) snapshot_download( repo_id=repo_id, local_dir=str(target_dir), local_dir_use_symlinks=False, resume_download=True, ) logger.success("✓ Download completed") except Exception as e: logger.error(f"✗ Download failed: {e}") logger.info( f"\nIf you are offline, please manually download from:\n" f" https://huggingface.co/{repo_id}\n" f"and place in: {target_dir.absolute()}" ) raise RuntimeError(f"Failed to download model weights: {e}") # Verify after download if manifest: all_valid, missing_after, corrupted_after = verify_file_integrity( target_dir, manifest, verify_checksums=verify ) if not all_valid: error_msg = [] if missing_after: error_msg.append(f"Still missing {len(missing_after)} file(s)") if corrupted_after: error_msg.append(f"Still corrupted {len(corrupted_after)} file(s)") raise FileNotFoundError( f"After download: {', '.join(error_msg)}\n" f"Please verify the download or manually place files in:\n" f" {target_dir.absolute()}" ) logger.success("✓ All model weights validated successfully") return target_dir