avans06/SeedVR2_Image_upscaler
8
1#!/usr/bin/env python32"""3Gradio front-end wrapper for SeedVR2's official inference_cli.py4 5This script is the user's app.py enhanced to stream subprocess logs in real-time6into the Gradio logs textbox. It runs the CLI as subprocesses and streams7stdout/stderr lines as they arrive using a queue and reader threads. The Gradio8handler `ui_upscale` is implemented as a generator so the frontend receives9incremental updates.10 11This script provides a simple web UI for single-image upscaling using the12official ComfyUI-SeedVR2_VideoUpscaler `inference_cli.py` script. It calls the13official CLI as a subprocess, and will automatically download model weights14from Hugging Face (numz/SeedVR2_comfyUI) if they are missing. If the15ComfyUI-SeedVR2_VideoUpscaler repository is not present, the script will16attempt to `git clone` it automatically into ./ComfyUI-SeedVR2_VideoUpscaler.17 18Run:19 python app.py20 21Requirements22- Python 3.10+23- Gradio (pip install gradio)24- Git available in PATH (for automatic cloning) or clone the repo manually25- PyTorch + CUDA (if using GPU)26 27Notes28- This wrapper calls the repo's `inference_cli.py` as a subprocess so the CLI's29 memory/optimization features (BlockSwap, VAE tiling, etc.) remain available.30- Models will be downloaded to the cloned repo's ./models/SeedVR2 directory if31 missing. Use HUGGINGFACE_HUB_TOKEN env var if required for private access.32"""33import os34import sys35import cv236import time37import torch38import queue39import shutil40import zipfile41import threading42import subprocess43import numpy as np44import gradio as gr45 46from pathlib import Path47from typing import Optional, Tuple, Generator, List48 49# huggingface helper (used for model auto-download)50from huggingface_hub import hf_hub_download51 52 53def imreadUTF8(path, flags=cv2.IMREAD_COLOR):54 """55 OpenCV's cv2.imread cannot handle non-ASCII paths.56 This function reads an image from a path that may contain UTF-8 characters.57 """58 try:59 # Use NumPy to read from the file, which correctly handles UTF-8 paths60 with open(path, "rb") as stream:61 bytes_data = bytearray(stream.read())62 numpyarray = np.asarray(bytes_data, dtype=np.uint8)63 # Use cv2.imdecode to decode the image from the memory buffer64 img = cv2.imdecode(numpyarray, flags)65 return img66 except Exception as e:67 # If reading fails, print the error message and return None68 print(f"ERROR: Failed to read image with UTF-8 path: {path}")69 print(f" Details: {e}")70 return None71 72def imwriteUTF8(save_path, image):73 """74 OpenCV's cv2.imwrite cannot handle non-ASCII paths.75 This function writes an image to a path that may contain UTF-8 characters.76 """77 try:78 img_name = os.path.basename(save_path)79 _, extension = os.path.splitext(img_name)80 # Encode the image into the specified format (determined by the file extension)81 is_success, im_buf_arr = cv2.imencode(extension, image)82 if is_success:83 # Write the image data from memory to the file84 im_buf_arr.tofile(save_path)85 return True86 else:87 print(f"ERROR: Failed to encode image for path: {save_path}")88 return False89 except Exception as e:90 print(f"ERROR: Failed to write image with UTF-8 path: {save_path}")91 print(f" Details: {e}")92 return False93 94# Apply Monkey Patch to cv2 (for app.py usage)95print("[SeedVR2 Gradio] Applying UTF-8 patch to OpenCV (Frontend)...")96cv2.imread = imreadUTF897cv2.imwrite = imwriteUTF898 99# ----------------100# Config / paths101# ----------------102REPO_URL = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler.git"103CLONE_DIR = Path(__file__).resolve().parent / "ComfyUI-SeedVR2_VideoUpscaler"104INFERENCE_CLI = CLONE_DIR / "inference_cli.py"105PY_EXE = sys.executable # Use same Python executable to run CLI106 107# Path to the custom improved blockswap file108IMPROVED_BLOCKSWAP_SOURCE = Path(__file__).resolve().parent / "src" / "optimization" / "blockswap.py"109IMPROVED_MEMORY_MANAGER_SOURCE = Path(__file__).resolve().parent / "src" / "optimization" / "memory_manager.py"110 111# Default HF repo for VAE (VAE is usually static and comes from the official repo)112DEFAULT_VAE_REPO_ID = "numz/SeedVR2_comfyUI"113# Models are now stored in a fixed top-level directory, independent of the clone dir114DEFAULT_MODEL_DIR = Path(__file__).resolve().parent / "models" / "SeedVR2"115 116# ----------------117# Model Definitions (RepoID / Filename)118# ----------------119# Standard Models (Safetensors)120MODEL_CHOICES = [121 "numz/SeedVR2_comfyUI/seedvr2_ema_3b_fp8_e4m3fn.safetensors",122 "numz/SeedVR2_comfyUI/seedvr2_ema_3b_fp16.safetensors",123 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_7b_fp8_e4m3fn_mixed_block35_fp16.safetensors",124 "numz/SeedVR2_comfyUI/seedvr2_ema_7b_fp16.safetensors",125 # sharp variants126 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_7b_sharp_fp8_e4m3fn_mixed_block35_fp16.safetensors",127 "numz/SeedVR2_comfyUI/seedvr2_ema_7b_sharp_fp16.safetensors",128]129 130# GGUF / alternate model support131GGUF_CHOICES = [132 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_3b-Q4_K_M.gguf",133 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_3b-Q8_0.gguf",134 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_7b-Q4_K_M.gguf",135 # sharp variants136 "AInVFX/SeedVR2_comfyUI/seedvr2_ema_7b_sharp-Q4_K_M.gguf",137 # custom GGUF from cmeka138 "cmeka/SeedVR2-GGUF/seedvr2_ema_7b-Q8_0.gguf",139 "cmeka/SeedVR2-GGUF/seedvr2_ema_7b_sharp-Q8_0.gguf",140]141 142# # Model registry with metadata143# MODEL_REGISTRY = {144# # 3B models145# "seedvr2_ema_3b-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="3B", precision="Q4_K_M", sha256="e665e3909de1a8c88a69c609bca9d43ff5a134647face2ce4497640cc3597f0e"),146# "seedvr2_ema_3b-Q8_0.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="3B", precision="Q8_0", sha256="be0d60083a2051a265eb4b77f28edf494e6db67ffc250216f32b72292e5cbd96"),147# "seedvr2_ema_3b_fp8_e4m3fn.safetensors": ModelInfo(size="3B", precision="fp8_e4m3fn", sha256="3bf1e43ebedd570e7e7a0b1b60d6a02e105978f505c8128a241cde99a8240cff"),148# "seedvr2_ema_3b_fp16.safetensors": ModelInfo(size="3B", precision="fp16", sha256="2fd0e03a3dad24e07086750360727ca437de4ecd456f769856e960ae93e2b304"),149 150# # 7B models151# "seedvr2_ema_7b-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="Q4_K_M", sha256="db9cb2ad90ebd40d2e8c29da2b3fc6fd03ba87cd58cbadceccca13ad27162789"),152# "seedvr2_ema_7b_fp8_e4m3fn_mixed_block35_fp16.safetensors": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="fp8_e4m3fn_mixed_block35_fp16", sha256="3d68b5ec0b295ae28092e355c8cad870edd00b817b26587d0cb8f9dd2df19bb2"),153# "seedvr2_ema_7b_fp16.safetensors": ModelInfo(size="7B", precision="fp16", sha256="7b8241aa957606ab6cfb66edabc96d43234f9819c5392b44d2492d9f0b0bbe4a"),154 155# # 7B sharp variants156# "seedvr2_ema_7b_sharp-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="Q4_K_M", variant="sharp", sha256="7aed800ac4eb8e0d18569a954c0ff35f5a1caa3ed5d920e66cc31405f75b6e69"),157# "seedvr2_ema_7b_sharp_fp8_e4m3fn_mixed_block35_fp16.safetensors": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="fp8_e4m3fn_mixed_block35_fp16", variant="sharp", sha256="0d2c5b8be0fda94351149c5115da26aef4f4932a7a2a928c6f184dda9186e0be"),158# "seedvr2_ema_7b_sharp_fp16.safetensors": ModelInfo(size="7B", precision="fp16", variant="sharp", sha256="20a93e01ff24beaeebc5de4e4e5be924359606c356c9c51509fba245bd2d77dd"),159 160# # VAE models161# "ema_vae_fp16.safetensors": ModelInfo(category="vae", precision="fp16", sha256="20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1"),162# }163 164# Detect Hardware Availability165CUDA_AVAILABLE = torch.cuda.is_available()166# Detect MPS availability (for Apple Silicon)167MPS_AVAILABLE = hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and torch.backends.mps.is_built()168# Check for any hardware acceleration169ACCELERATOR_AVAILABLE = CUDA_AVAILABLE or MPS_AVAILABLE170 171# -----------------172# Repo / model helpers173# -----------------174def ensure_repo_cloned(175 repo_url: str = REPO_URL, 176 clone_dir_name: str = "ComfyUI-SeedVR2_VideoUpscaler", 177 repo_branch: str = "",178 force_update: bool = False179) -> Path:180 """181 Ensure the repository is cloned locally into the specified directory name.182 Supports specific branches/tags via checkout.183 Returns the resolved Path object to the cloned directory.184 """185 # Resolve the physical path based on the script's parent location186 target_clone_dir = Path(__file__).resolve().parent / clone_dir_name187 target_cli = target_clone_dir / "inference_cli.py"188 189 # Helper function to handle detached/orphaned commits190 def _smart_checkout(cwd, ref):191 print(f"[SeedVR2 Gradio] Checking out '{ref}' in {cwd} ...")192 try:193 # Try standard checkout first (fastest if ref exists locally)194 subprocess.run(["git", "-C", str(cwd), "checkout", ref], check=True)195 except subprocess.CalledProcessError:196 # Fallback: If ref is not found (e.g. orphaned commit hash), fetch it explicitly197 print(f"[SeedVR2 Gradio] Standard checkout failed. Attempting to fetch specific ref '{ref}' from origin...")198 try:199 subprocess.run(["git", "-C", str(cwd), "fetch", "origin", ref], check=True)200 subprocess.run(["git", "-C", str(cwd), "checkout", ref], check=True)201 except Exception as e:202 raise RuntimeError(f"Failed to fetch/checkout specific ref '{ref}': {e}")203 204 if target_clone_dir.exists() and (target_clone_dir / ".git").exists():205 # Repo exists206 if force_update:207 try:208 print(f"[SeedVR2 Gradio] Updating {target_clone_dir} ...")209 subprocess.run(["git", "-C", str(target_clone_dir), "fetch", "--all"], check=True)210 211 # If a specific branch/hash is requested212 if repo_branch:213 _smart_checkout(target_clone_dir, repo_branch)214 # If it's a branch name (not a detached hash), we might want to pull latest215 # But checking if it's a branch vs hash is complex, generally strictly checking out the ref is safer for reproducibility216 else:217 subprocess.run(["git", "-C", str(target_clone_dir), "pull"], check=True)218 except Exception as e:219 raise RuntimeError(f"Failed to update repository {target_clone_dir}: {e}")220 221 # If not forcing update, but a branch is specified, ensure we are on it222 elif repo_branch:223 try:224 subprocess.run(["git", "-C", str(target_clone_dir), "fetch", "--all"], check=True)225 _smart_checkout(target_clone_dir, repo_branch)226 except Exception as e:227 raise RuntimeError(f"Failed to switch to branch {repo_branch}: {e}")228 229 # Ensure inference_cli present230 if not target_cli.exists():231 raise RuntimeError(f"Repository found at {target_clone_dir} but inference_cli.py is missing.")232 233 return target_clone_dir234 235 # Clone repo if not exists236 try:237 print(f"[SeedVR2 Gradio] Cloning {repo_url} into {target_clone_dir} ...")238 239 # Standard clone (fetches default branch)240 subprocess.run(["git", "clone", repo_url, str(target_clone_dir)], check=True)241 242 if repo_branch:243 _smart_checkout(target_clone_dir, repo_branch)244 245 except FileNotFoundError:246 raise RuntimeError("git not found: please install Git or clone the repository manually.")247 except Exception as e:248 raise RuntimeError(f"Failed to clone repository: {e}")249 250 if not target_cli.exists():251 raise RuntimeError(f"Clone completed but inference_cli.py not found in {target_clone_dir}.")252 253 return target_clone_dir254 255 256def apply_inference_cli_patch(cli_path: Path):257 """258 Injects UTF-8 compatible imread/imwrite wrappers directly into inference_cli.py.259 This modifies the physical file so the subprocess (even on Windows spawn) uses the patch.260 """261 if not cli_path.exists():262 return263 264 try:265 with open(cli_path, "r", encoding="utf-8") as f:266 content = f.read()267 268 # Check if already patched to avoid duplicates269 if "def imreadUTF8" in content:270 return271 272 # The patch content to inject. 273 # Note: We ensure 'import numpy as np' and 'import os' are available or re-imported.274 # inference_cli.py typically has 'import cv2', we inject right after that.275 patch_code = r'''276# =============================================================================277# GRADIO APP PATCH: UTF-8 Support for Windows (Auto-Injected)278# =============================================================================279import numpy as np280import os281 282def imreadUTF8(path, flags=cv2.IMREAD_COLOR):283 try:284 with open(path, "rb") as stream:285 bytes_data = bytearray(stream.read())286 numpyarray = np.asarray(bytes_data, dtype=np.uint8)287 return cv2.imdecode(numpyarray, flags)288 except Exception as e:289 print(f"Error reading image {path}: {e}")290 return None291 292def imwriteUTF8(save_path, image):293 try:294 img_name = os.path.basename(save_path)295 _, extension = os.path.splitext(img_name)296 is_success, im_buf_arr = cv2.imencode(extension, image)297 if is_success:298 im_buf_arr.tofile(save_path)299 return True300 else:301 return False302 except Exception as e:303 print(f"Error writing image {save_path}: {e}")304 return False305 306# Override cv2 methods307cv2.imread = imreadUTF8308cv2.imwrite = imwriteUTF8309# =============================================================================310'''311 # Inject after 'import cv2'312 if "import cv2" in content:313 print(f"[SeedVR2 Gradio] Patching {cli_path} for UTF-8 subprocess support...")314 new_content = content.replace("import cv2", "import cv2" + patch_code, 1)315 with open(cli_path, "w", encoding="utf-8") as f:316 f.write(new_content)317 else:318 print("[SeedVR2 Gradio] WARNING: Could not find 'import cv2' in inference_cli.py. UTF-8 patch skipped.")319 320 except Exception as e:321 print(f"[SeedVR2 Gradio] ERROR applying UTF-8 patch to inference_cli: {e}")322 323 324def patch_model_registry(repo_root: Path):325 """326 Appends custom model definitions to src/utils/model_registry.py.327 This allows the CLI to recognize new GGUF models that aren't in the official registry.328 """329 registry_path = repo_root / "src" / "utils" / "model_registry.py"330 331 if not registry_path.exists():332 print(f"[SeedVR2 Gradio] WARN: Could not find model_registry.py at {registry_path}")333 return334 335 try:336 with open(registry_path, "r", encoding="utf-8") as f:337 content = f.read()338 339 # Check if already patched340 if "seedvr2_ema_7b-Q8_0.gguf" in content:341 return342 343 print(f"[SeedVR2 Gradio] Patching {registry_path} with custom GGUF models...")344 345 # Code to append to the end of the file. 346 # Since ModelInfo and MODEL_REGISTRY are defined in the file, we can use them directly.347 patch_code = r'''348 349# =============================================================================350# GRADIO APP PATCH: Custom Model Registry Entries351# =============================================================================352try:353 # Update registry with custom GGUF models requested by user354 MODEL_REGISTRY.update({355 "seedvr2_ema_7b-Q8_0.gguf": ModelInfo(repo="cmeka/SeedVR2-GGUF", size="7B", precision="Q8_0", sha256="669788655e8f15f306284f267a444e9766c8a421869577b16a961e43029c737b"),356 "seedvr2_ema_7b_sharp-Q8_0.gguf": ModelInfo(repo="cmeka/SeedVR2-GGUF", size="7B", precision="Q8_0", variant="sharp", sha256="b1f81cb5700b0b1f432f2c785528356c952c41c74d03d205c6f14b0bd6da303d"),357 })358 print("[Internal] Custom GGUF models injected into MODEL_REGISTRY successfully.")359except Exception as e:360 print(f"[Internal] Failed to inject custom models: {e}")361# =============================================================================362'''363 with open(registry_path, "a", encoding="utf-8") as f:364 f.write(patch_code)365 366 except Exception as e:367 print(f"[SeedVR2 Gradio] ERROR patching model_registry.py: {e}")368 369 370# -----------------371# BlockSwap Management372# -----------------373def manage_blockswap_file(use_improved: bool, repo_root: Path) -> str:374 """375 Manages the blockswap.py file in the specified cloned repository.376 Accepted `repo_root` path to ensure we modify the correct repo.377 """378 target_path_blockswap = repo_root / "src" / "optimization" / "blockswap.py"379 backup_path_blockswap = repo_root / "src" / "optimization" / "blockswap.py.bak"380 target_path_memory_manager = repo_root / "src" / "optimization" / "memory_manager.py"381 backup_path_memory_manager = repo_root / "src" / "optimization" / "memory_manager.py.bak"382 383 msg = ""384 # Ensure src/optimization exists (some forks might differ in structure)385 if not target_path_blockswap.parent.exists():386 return f"[WARN] Optimization folder not found at {target_path_blockswap.parent}. Skipping blockswap patch.\n"387 388 if use_improved:389 if not IMPROVED_BLOCKSWAP_SOURCE.exists():390 return f"[WARN] Improved blockswap source not found at {IMPROVED_BLOCKSWAP_SOURCE}. Keeping current version.\n"391 392 # 1. Check if we need to backup the original blockswap (only if backup doesn't exist yet)393 if target_path_blockswap.exists() and not backup_path_blockswap.exists():394 try:395 shutil.move(str(target_path_blockswap), str(backup_path_blockswap))396 msg += f"[INFO] Backed up original blockswap to {backup_path_blockswap.name}.\n"397 except Exception as e:398 return f"[ERROR] Failed to backup blockswap: {e}\n"399 400 # 2. Copy the improved file to target blockswap401 try:402 shutil.copy(str(IMPROVED_BLOCKSWAP_SOURCE), str(target_path_blockswap))403 msg += "[INFO] Switched to Improved BlockSwap (Nunchaku implementation).\n"404 except Exception as e:405 return f"[ERROR] Failed to install improved blockswap: {e}\n"406 407 # Memory Manager Handling408 if not IMPROVED_MEMORY_MANAGER_SOURCE.exists():409 return f"[WARN] Improved memory_manager source not found at {IMPROVED_MEMORY_MANAGER_SOURCE}. Keeping current version.\n"410 411 # 3. Check if we need to backup the original memory_manager (only if backup doesn't exist yet)412 if target_path_memory_manager.exists() and not backup_path_memory_manager.exists():413 try:414 shutil.move(str(target_path_memory_manager), str(backup_path_memory_manager))415 msg += f"[INFO] Backed up original memory_manager to {backup_path_memory_manager.name}.\n"416 except Exception as e:417 return f"[ERROR] Failed to backup memory_manager: {e}\n"418 419 # 4. Copy the improved file to target memory_manager420 try:421 shutil.copy(str(IMPROVED_MEMORY_MANAGER_SOURCE), str(target_path_memory_manager))422 msg += "[INFO] Switched to Improved memory_manager (Nunchaku implementation).\n"423 except Exception as e:424 return f"[ERROR] Failed to install improved memory_manager: {e}\n"425 426 return msg427 428 else:429 # Restore original blockswap if available430 if backup_path_blockswap.exists():431 try:432 # Remove current target blockswap if it exists (which might be the improved one)433 if target_path_blockswap.exists():434 os.remove(target_path_blockswap)435 436 # Restore backup437 shutil.move(str(backup_path_blockswap), str(target_path_blockswap))438 msg += "[INFO] Restored Original BlockSwap from backup.\n"439 except Exception as e:440 return f"[ERROR] Failed to restore original blockswap: {e}\n"441 else:442 # Backup doesn't exist, assume we are already on original or clean install443 msg += "[INFO] Using Original BlockSwap (No backup found/needed).\n"444 445 # Restore original memory_manager if available446 if backup_path_memory_manager.exists():447 try:448 # Remove current target memory_manager if it exists (which might be the improved one)449 if target_path_memory_manager.exists():450 os.remove(target_path_memory_manager)451 452 # Restore backup453 shutil.move(str(backup_path_memory_manager), str(target_path_memory_manager))454 msg += "[INFO] Restored Original memory_manager from backup.\n"455 except Exception as e:456 return f"[ERROR] Failed to restore original memory_manager: {e}\n"457 else:458 # Backup doesn't exist, assume we are already on original or clean install459 msg += "[INFO] Using Original memory_manager (No backup found/needed).\n"460 461 return msg462 463 464# -----------------465# Model download466# -----------------467def ensure_models_available(468 selected_model_filename: str, 469 model_dir: Optional[Path] = None, 470 repo_id: str = DEFAULT_VAE_REPO_ID471) -> None:472 """473 Ensure the selected DiT model and the VAE file exist locally.474 475 If missing, download from the specified Hugging Face repo directly into model_dir476 using 'local_dir' to avoid nested cache structures.477 """478 479 if model_dir is None:480 model_dir = DEFAULT_MODEL_DIR481 else:482 model_dir = Path(model_dir)483 484 model_dir.mkdir(parents=True, exist_ok=True)485 486 # Items to check: VAE + selected DiT model487 required = ["ema_vae_fp16.safetensors", selected_model_filename]488 489 # Check if files physically exist at the target location490 missing = [_f for _f in required if not (model_dir / _f).exists()]491 492 if not missing:493 return494 495 # Optional: silence HF symlink warning if desired496 os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")497 498 # Attempt download for each missing file499 hf_token = os.environ.get("HF_ACCESS_TOKEN")500 501 for fname in missing:502 target_path = model_dir / fname503 # If file already somehow exists at target, skip504 if target_path.exists():505 continue506 507 # Decide repo for this filename:508 # - VAE must always come from the official DEFAULT_VAE_REPO_ID (numz/SeedVR2_comfyUI)509 # - Dit model uses the provided repo_id (which comes from the dropdown selection)510 repo_for_fname = DEFAULT_VAE_REPO_ID if fname == "ema_vae_fp16.safetensors" else repo_id511 512 try:513 print(f"[SeedVR2 Gradio] Downloading {fname} from {repo_for_fname} directly to {model_dir} ...")514 515 # Use local_dir instead of cache_dir.516 # This forces the file to be saved exactly at {model_dir}/{fname}517 # local_dir_use_symlinks=False ensures we get a real file, not a symlink,518 # which prevents issues where the CLI subprocess cannot resolve the path.519 downloaded_path = hf_hub_download(520 repo_id=repo_for_fname,521 filename=fname,522 local_dir=str(model_dir), # Download directly to target folder523 repo_type="model",524 token=hf_token,525 )526 print(f"[SeedVR2 Gradio] Download completed: {downloaded_path}")527 528 except Exception as e:529 raise RuntimeError(f"Failed to download {fname} from Hugging Face repo {repo_for_fname}: {e}")530 531 532# ----------------533# Subprocess streaming helpers534# ----------------535def _start_process_stream(cmd_args, cwd: str, env: dict) -> Tuple[Optional[subprocess.Popen], queue.Queue, Optional[threading.Thread], Optional[threading.Thread]]:536 """Start subprocess and return (proc, q, t_out, t_err).537 538 The returned queue will receive text lines as they arrive. Lines are simple539 strings (already newline-terminated). stderr lines are prefixed with "stderr: ".540 """541 q = queue.Queue()542 543 try:544 proc = subprocess.Popen(545 cmd_args,546 cwd=cwd,547 stdout=subprocess.PIPE,548 stderr=subprocess.PIPE,549 text=True,550 encoding='utf-8',551 errors='replace',552 bufsize=1,553 env=env554 )555 except Exception as e:556 # Put error to queue and return a dummy proc557 q.put(f"[FAILED TO LAUNCH] {e}\n")558 return None, q, None, None559 560 def _reader(fh, prefix: str):561 try:562 while True:563 line = fh.readline()564 if not line:565 break566 if not line.endswith("\n"):567 line = line + "\n"568 q.put(prefix + line)569 except Exception as e:570 q.put(f"[reader error] {e}\n")571 572 t_out = threading.Thread(target=_reader, args=(proc.stdout, ""), daemon=True)573 t_err = threading.Thread(target=_reader, args=(proc.stderr, "stderr: "), daemon=True)574 t_out.start()575 t_err.start()576 return proc, q, t_out, t_err577 578 579# -----------------580# CLI runner (streaming)581# -----------------582def expected_upscaled_path(input_path: str, output_format: str = "png") -> str:583 """Calculates the expected output path based on the input path and requested format."""584 p = Path(input_path)585 stem = p.stem586 parent = p.parent587 suffix = "_upscaled"588 # if output_format == "mp4":589 # return str((parent / f"{stem}{suffix}.mp4").resolve())590 # else:591 # return str((parent / f"{stem}{suffix}.png").resolve())592 593 return str((parent / f"{stem}{suffix}.{output_format}").resolve())594 595 596# Single-image/video CLI runner (generator)597def run_cli_upscale_stream(598 input_path: str,599 resolution: int = 1080,600 max_resolution: int = 0,601 dit_model_filename: Optional[str] = None, # Receives just the filename602 cuda_device: Optional[str] = None,603 # Compilation & Performance604 compile_dit: bool = False,605 compile_vae: bool = False,606 compile_backend: str = "inductor",607 compile_mode: str = "default",608 compile_fullgraph: bool = False,609 compile_dynamic: bool = False,610 compile_dynamo_cache_size_limit: int = 64,611 compile_dynamo_recompile_limit: int = 128,612 attention_mode: str = "sdpa",613 # Tiling (Split Encode/Decode)614 vae_encode_tiled: bool = False,615 vae_encode_tile_size: int = 1024,616 vae_encode_tile_overlap: int = 128,617 vae_decode_tiled: bool = False,618 vae_decode_tile_size: int = 1024,619 vae_decode_tile_overlap: int = 128,620 tile_debug: str = "false",621 # Processing622 batch_size: int = 1,623 uniform_batch_size: bool = False,624 seed: int = 42,625 skip_first_frames: int = 0,626 load_cap: int = 0,627 # Quality & Color628 color_correction: str = "lab",629 input_noise_scale: float = 0.0,630 latent_noise_scale: float = 0.0,631 # Memory & Offload632 blocks_to_swap: int = 0,633 swap_io_components: bool = False,634 dit_offload_device: str = "none",635 vae_offload_device: str = "none",636 tensor_offload_device: str = "cpu",637 cache_dit: bool = False,638 cache_vae: bool = False,639 extra_args: str = "",640 model_dir: Optional[str] = None,641 repo_id: str = DEFAULT_VAE_REPO_ID, # Receives the specific Repo ID for DiT642 repo_path: Optional[Path] = None,643 timeout: int = 3600,644 pre_downscale: bool = False, # for artifact removal645 downscale_rate: float = 0.5,646 output_format: str = "png", # Can now be "mp4"647 use_improved_blockswap: bool = False, # New argument for switching blockswap version648 # Video Args649 chunk_size: int = 0,650 temporal_overlap: int = 0,651 prepend_frames: int = 0,652 video_backend: str = "opencv",653 use_10bit: bool = False,654 # Debug Arg655 debug: bool = False656) -> Generator[Tuple[Optional[str], str], None, None]:657 """658 Generator yields (out_path_or_None, logs_so_far) while streaming CLI logs.659 Includes Phase-Aware Dynamic Fallback logic.660 """661 # Defaults662 if repo_path is None:663 repo_path = CLONE_DIR664 665 current_inference_cli = repo_path / "inference_cli.py"666 667 # 1. Repo Check668 if not current_inference_cli.exists():669 yield None, f"[ERROR] inference_cli.py not found in {repo_path}\n"670 return671 672 # Patch inference_cli.py with UTF-8 support673 try:674 apply_inference_cli_patch(current_inference_cli)675 except Exception as e:676 yield None, f"[WARN] Failed to patch inference_cli: {e}\n"677 678 679 # Patch model_registry.py with custom models680 try:681 patch_model_registry(repo_path)682 except Exception as e:683 yield None, f"[WARN] Failed to patch model_registry: {e}\n"684 685 # Handle BlockSwap File Replacement Logic686 try:687 swap_log = manage_blockswap_file(use_improved_blockswap, repo_root=repo_path)688 # Yield the log about blockswap immediately689 yield None, swap_log690 except Exception as e:691 yield None, f"[ERROR] BlockSwap management failed: {e}\n"692 693 # Use the global default if not provided694 if model_dir is None:695 model_dir = str(DEFAULT_MODEL_DIR)696 697 # Ensure model files present698 if dit_model_filename:699 try:700 ensure_models_available(701 dit_model_filename, 702 model_dir=Path(model_dir), 703 repo_id=repo_id,704 )705 except Exception as e:706 yield None, f"[ERROR] Model download failed: {e}\n"707 return708 709 safe_input_path = input_path710 temp_copy = None711 712 # Pre-downscale logic (Artifact Removal Trick) - Only applies to Images in this implementation713 # We skip this for MP4 files to avoid complex video processing in python before CLI714 is_video = input_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))715 716 # Pre-downscale (Images only)717 if pre_downscale and not is_video:718 try:719 filename = os.path.basename(input_path)720 # Load original image using OpenCV721 img_obj = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)722 if img_obj is None:723 raise ValueError(f"Failed to load image: {input_path}")724 725 # Calculate new dimensions (OpenCV shape is [height, width])726 h, w = img_obj.shape[:2]727 if (max(w, h) > 250):728 new_w = int(w * downscale_rate)729 new_h = int(h * downscale_rate)730 731 # Resize732 # Use INTER_AREA for downscaling (better quality/less aliasing for shrinking)733 # Use INTER_LANCZOS4 if scaling up (though this block is specifically for downscaling)734 interpolation_method = cv2.INTER_AREA if downscale_rate < 1.0 else cv2.INTER_LANCZOS4735 736 img_resized = cv2.resize(img_obj, (new_w, new_h), interpolation=interpolation_method)737 738 # Prepare temp directory739 tmp_dir = CLONE_DIR / "tmp_inputs"740 tmp_dir.mkdir(parents=True, exist_ok=True)741 742 # Save to a unique temp file (forces .png for intermediate input)743 new_name = f"{filename}_downscaled.png"744 temp_copy = str(tmp_dir / new_name)745 # Use patched cv2.imwrite746 cv2.imwrite(temp_copy, img_resized)747 748 # Use this temp file as the input for CLI749 safe_input_path = temp_copy750 751 yield None, f"[INFO] Pre-downscaled input by factor {downscale_rate} (Size: {w}x{h} -> {new_w}x{new_h}) to reduce artifacts.\n"752 753 except Exception as e:754 yield None, f"[ERROR] Failed to pre-downscale image: {e}\n"755 return756 757 # 2. Command Builder758 def _build_cmd(curr_tile_size, curr_batch_size):759 # Determine strict output format760 if is_video:761 # If input is video, force mp4 output for CLI unless user explicitly wants png sequence?762 # Usually users want mp4 back. 763 cmd_format = "mp4"764 else:765 # For images, use png (CLI handles webp/jpg conversion internally if modified, 766 # but standard CLI outputs png/mp4). We force png here, app.py handles conversion later.767 cmd_format = "png"768 769 cmd = [PY_EXE, str(current_inference_cli), safe_input_path,770 "--resolution", str(resolution),771 "--output_format", cmd_format, 772 "--batch_size", str(curr_batch_size),773 "--color_correction", color_correction,774 "--model_dir", str(model_dir),775 "--seed", str(seed),776 "--attention_mode", str(attention_mode)]777 778 if max_resolution and int(max_resolution) > 0:779 cmd += ["--max_resolution", str(int(max_resolution))]780 781 if dit_model_filename:782 # CLI just needs the filename relative to --model_dir (or absolute path)783 cmd += ["--dit_model", str(dit_model_filename)]784 785 # Only add --cuda_device if CUDA available and user provided a value786 if CUDA_AVAILABLE and cuda_device:787 cmd += ["--cuda_device", str(cuda_device)]788 789 # --- Compilation Options ---790 if compile_dit:791 cmd += ["--compile_dit"]792 793 if compile_vae:794 cmd += ["--compile_vae"]795 796 if compile_dit or compile_vae:797 cmd += [798 "--compile_backend", str(compile_backend),799 "--compile_mode", str(compile_mode),800 "--compile_dynamo_cache_size_limit", str(compile_dynamo_cache_size_limit),801 "--compile_dynamo_recompile_limit", str(compile_dynamo_recompile_limit)802 ]803 if compile_fullgraph:804 cmd += ["--compile_fullgraph"]805 if compile_dynamic:806 cmd += ["--compile_dynamic"]807 808 # --- Tiling Options ---809 # Note: curr_tile_size comes from the loop strategy (Phase Fallback), 810 # normally we use the user provided vae_encode_tile_size.811 812 if vae_encode_tiled:813 cmd += ["--vae_encode_tiled", 814 "--vae_encode_tile_size", str(curr_tile_size),815 "--vae_encode_tile_overlap", str(vae_encode_tile_overlap)]816 817 if vae_decode_tiled:818 cmd += ["--vae_decode_tiled", 819 "--vae_decode_tile_size", str(vae_decode_tile_size), 820 "--vae_decode_tile_overlap", str(vae_decode_tile_overlap)]821 822 if tile_debug != "false":823 cmd += ["--tile_debug", str(tile_debug)]824 825 # --- Processing & Quality ---826 if uniform_batch_size:827 cmd += ["--uniform_batch_size"]828 829 if skip_first_frames > 0:830 cmd += ["--skip_first_frames", str(int(skip_first_frames))]831 832 if load_cap > 0:833 cmd += ["--load_cap", str(int(load_cap))]834 835 if input_noise_scale > 0:836 cmd += ["--input_noise_scale", str(input_noise_scale)]837 838 if latent_noise_scale > 0:839 cmd += ["--latent_noise_scale", str(latent_noise_scale)]840 841 # --- BlockSwap / Offload / Caching ---842 if blocks_to_swap and int(blocks_to_swap) > 0:843 cmd += ["--blocks_to_swap", str(int(blocks_to_swap))]844 845 if swap_io_components:846 cmd += ["--swap_io_components"]847 848 # Offload flags: note these are strings like "none"/"cpu"/"cuda:0"849 if dit_offload_device and dit_offload_device != "none":850 # Ensure we don't pass a cuda device offload when cuda isn't available851 if not (dit_offload_device.startswith("cuda") and not CUDA_AVAILABLE):852 cmd += ["--dit_offload_device", str(dit_offload_device)]853 854 if vae_offload_device and vae_offload_device != "none":855 if not (vae_offload_device.startswith("cuda") and not CUDA_AVAILABLE):856 cmd += ["--vae_offload_device", str(vae_offload_device)]857 858 if tensor_offload_device and tensor_offload_device != "none":859 if not (tensor_offload_device.startswith("cuda") and not CUDA_AVAILABLE):860 cmd += ["--tensor_offload_device", str(tensor_offload_device)]861 862 if cache_dit:863 cmd += ["--cache_dit"]864 if cache_vae:865 cmd += ["--cache_vae"]866 867 # --- Video Specific Flags ---868 if chunk_size > 0:869 cmd += ["--chunk_size", str(int(chunk_size))]870 871 if temporal_overlap > 0:872 cmd += ["--temporal_overlap", str(int(temporal_overlap))]873 874 if prepend_frames > 0:875 cmd += ["--prepend_frames", str(int(prepend_frames))]876 877 if video_backend and video_backend != "opencv":878 cmd += ["--video_backend", str(video_backend)]879 880 if use_10bit:881 cmd += ["--10bit"]882 883 # Debug Flag884 if debug:885 cmd += ["--debug"]886 887 if extra_args:888 # Allow advanced users to type additional flags (space separated)889 cmd += extra_args.split()890 891 return cmd892 893 # 3. Dynamic Strategy Loop894 # Use encode tile size as the dynamic variable for fallback895 current_tile_size = int(vae_encode_tile_size)896 current_batch_size = int(batch_size)897 898 # Initialize log tracking899 logs_buf = ""900 # Add previous swap logs to buf901 logs_buf += swap_log if 'swap_log' in locals() else ""902 903 max_attempts = 5 # Prevent infinite loops904 attempt_count = 0905 906 idx = 0907 while attempt_count < max_attempts:908 attempt_count += 1909 910 note = f"Tile: {current_tile_size}, Batch: {current_batch_size}"911 header = f"\n\n=== ATTEMPT {attempt_count}/{max_attempts} ({note}) ===\n"912 logs_buf += header913 # yield immediate header914 yield None, logs_buf915 916 cmd = _build_cmd(current_tile_size, current_batch_size)917 logs_buf += f"[CMD] {' '.join(cmd)}\n"918 yield None, logs_buf919 920 # start streaming process921 env = os.environ.copy()922 # Make Python in child process print using UTF-8 (avoids cp950 UnicodeEncodeError on Windows)923 env['PYTHONIOENCODING'] = 'utf-8'924 env['PYTHONUTF8'] = '1'925 # # help fragmentation/alloc issues; user may tune926 # env.setdefault('PYTORCH_ALLOC_CONF', os.environ.get('PYTORCH_ALLOC_CONF', 'max_split_size_mb:128'))927 928 proc, q, t_out, t_err = _start_process_stream(cmd, cwd=str(CLONE_DIR), env=env)929 if proc is None:930 logs_buf += "[ERROR] Failed to launch subprocess.\n"931 yield None, logs_buf932 break # try next strategy? here treat as fatal933 934 # State tracking for this run935 current_phase = "init" # init, vae_enc, dit, vae_dec, post936 oom_detected = False937 start_time = time.time()938 939 # poll queue940 while True:941 try:942 # wait up to 0.5s for a line943 line = q.get(timeout=0.5)944 logs_buf += line945 yield None, logs_buf946 947 lower_line = line.lower()948 949 # Track Phase950 if "phase 1: vae encoding" in lower_line:951 current_phase = "vae_enc"952 elif "phase 2: dit upscaling" in lower_line:953 current_phase = "dit"954 elif "phase 3: vae decoding" in lower_line:955 current_phase = "vae_dec"956 elif "saving" in lower_line or "converting" in lower_line:957 current_phase = "post"958 959 # Check for OOM960 oom_indicators = ["outofmemory", "out of memory", "allocation on device", "oom", "cuda out of memory"]961 if any(k in lower_line for k in oom_indicators):962 logs_buf += f"\n[WARN] OOM detected during phase: {current_phase.upper()}\n"963 yield None, logs_buf964 oom_detected = True965 try:966 proc.kill() # Kill immediately to recover VRAM967 except: pass968 break969 970 except queue.Empty:971 # no new line - check process status972 if proc.poll() is not None:973 break974 # still running - continue polling975 continue976 977 # Flush remaining978 while True:979 try:980 line = q.get_nowait()981 logs_buf += line982 yield None, logs_buf983 except queue.Empty:984 break985 986 # Wait for reader threads to exit987 try:988 if t_out:989 t_out.join(timeout=1)990 if t_err:991 t_err.join(timeout=1)992 except Exception:993 pass994 995 runtime = time.time() - start_time996 logs_buf += f"[Attempt {idx} finished in {runtime:.2f}s] returncode={proc.returncode}\n"997 idx += 1998 yield None, logs_buf999 1000 # 4. Success Check using safe_input_path1001 # CLI generates output relative to the actual input file used (which might be the temp one)1002 # For video, strict output detection logic1003 out_fmt_check = "mp4" if is_video else "png"1004 out_path = expected_upscaled_path(safe_input_path, output_format=out_fmt_check)1005 1006 if Path(out_path).exists():1007 logs_buf += f"[SUCCESS] Intermediate Output: {out_path}\n"1008 # Cleanup temp file if we created one1009 if temp_copy:1010 try:1011 os.remove(temp_copy)1012 except Exception:1013 pass1014 yield out_path, logs_buf1015 return1016 1017 # 5. Failure Analysis & Parameter Adjustment1018 if oom_detected or proc.returncode != 0:1019 logs_buf += f"\n[INFO] Attempt {attempt_count} failed. Analyzing OOM Phase: {current_phase.upper()}...\n"1020 1021 # --- INTELLIGENT ADJUSTMENT LOGIC ---1022 1023 # Case A: VAE OOM (Phase 1 or 3) -> Reduce Tile Size1024 if current_phase in ["vae_enc", "vae_dec"]:1025 if current_tile_size > 256:1026 new_tile = max(256, current_tile_size // 2)1027 logs_buf += f"[STRATEGY] VAE OOM detected. Reducing Tile Size: {current_tile_size} -> {new_tile}\n"1028 current_tile_size = new_tile1029 else:1030 # Tile size already min, try reducing batch size as a last resort1031 new_batch = max(1, current_batch_size // 2)1032 logs_buf += f"[STRATEGY] VAE OOM but Tile Size is min. Reducing Batch Size: {current_batch_size} -> {new_batch}\n"1033 current_batch_size = new_batch1034 1035 # Case B: DiT OOM (Phase 2) -> Reduce Batch Size1036 elif current_phase == "dit":1037 if current_batch_size > 1:1038 # For video consistency, try to keep 4n+1 if possible, or just halve it1039 new_batch = max(1, current_batch_size // 2)1040 logs_buf += f"[STRATEGY] DiT OOM detected. Reducing Batch Size: {current_batch_size} -> {new_batch}\n"1041 current_batch_size = new_batch1042 else:1043 logs_buf += f"[FAIL] DiT OOM with Batch Size 1. Cannot reduce further.\n"1044 break1045 1046 # Case C: Post-Process OOM (Phase 4) -> Reduce Batch Size1047 elif current_phase == "post":1048 if current_batch_size > 1:1049 new_batch = max(1, current_batch_size // 2)1050 logs_buf += f"[STRATEGY] Post-Process OOM detected. Reducing Batch Size: {current_batch_size} -> {new_batch}\n"1051 current_batch_size = new_batch1052 else:1053 logs_buf += "[FAIL] Post-Process OOM with Batch Size 1.\n"1054 break1055 1056 # Case D: Unknown/Init OOM -> Reduce both safely1057 else:1058 current_tile_size = max(256, current_tile_size // 2)1059 current_batch_size = max(1, current_batch_size // 2)1060 logs_buf += f"[STRATEGY] Early OOM. Reducing both Tile ({current_tile_size}) and Batch ({current_batch_size}).\n"1061 1062 # Check if we are just retrying same settings (infinite loop prevention)1063 if attempt_count >= max_attempts:1064 logs_buf += "[FAIL] Max attempts reached.\n"1065 break1066 1067 yield None, logs_buf1068 # Loop continues with new settings1069 else:1070 # Non-OOM fatal error1071 logs_buf += f"[ERROR] Non-zero return code (not OOM) - stopping.\n"1072 yield None, logs_buf1073 return1074 1075 # all strategies exhausted1076 logs_buf += "[FAILED] No output produced after all strategies.\n"1077 if temp_copy:1078 try:1079 os.remove(temp_copy)1080 except:1081 pass1082 yield None, logs_buf1083 return1084 1085 1086# --- Preset change handler (considers CUDA & MPS availability) ---1087def preset_changed(preset_value):1088 # Updated Tuple Order:1089 # 0: compile_dit, 1: compile_vae, 1090 # 2: vae_encode_tiled, 3: vae_encode_tile_size, 1091 # 4: vae_decode_tiled, 5: vae_decode_tile_size,1092 # 6: max_resolution, 7: blocks_to_swap, 8: swap_io_components1093 # 9: dit_offload_device, 10: vae_offload_device, 11: tensor_offload_device, 1094 # 12: extra_args, 13: chunk_size, 14: temporal_overlap1095 1096 if preset_value == "Recommended (low VRAM)":1097 return (1098 False, # compile_dit1099 False, # compile_vae1100 True, # vae_encode_tiled1101 512, # vae_encode_tile_size1102 True, # vae_decode_tiled1103 512, # vae_decode_tile_size (sync with encode for safety)1104 1920, # max_resolution1105 32, # blocks_to_swap1106 True, # swap_io_components1107 "cpu", # dit_offload_device1108 "none", # vae_offload_device (Keep VAE on device if possible)1109 "cpu", # tensor_offload_device (Offload tensors to save VRAM)1110 "--blocks_to_swap 0", # extra_args1111 0, # chunk_size1112 0 # temporal_overlap (0=auto/disabled)1113 )1114 elif preset_value == "Offload (very slow)":1115 return (1116 False, # compile_dit1117 False, # compile_vae1118 True, # vae_encode_tiled1119 256, # vae_encode_tile_size1120 True, # vae_decode_tiled1121 256, # vae_decode_tile_size1122 1440, # max_resolution1123 99, # blocks_to_swap1124 True, # swap_io_components1125 "cpu", # dit_offload_device1126 "cpu", # vae_offload_device1127 "cpu", # tensor_offload_device1128 "--blocks_to_swap 99 --swap_io_components --dit_offload_device cpu --vae_offload_device cpu --tensor_offload_device cpu",1129 0,1130 01131 )1132 elif preset_value == "High quality (fast if lots of VRAM)":1133 return (1134 True,1135 True,1136 False,1137 512, 1138 False,1139 512,1140 0,1141 0,1142 False,1143 "none",1144 "none",1145 "none", # Keep tensors on GPU/MPS1146 "",1147 0,1148 01149 )1150 # fallback1151 return (False, False, True, 256, True, 256, 1920, 0, False, "none", "none", "cpu", "--blocks_to_swap 0", 0, 0)1152 1153 1154# ---------------- Paste JS (attach to gallery elem) ----------------1155paste_js = """1156function initPaste() {1157 document.addEventListener('paste', function(e) {1158 const gallery = document.getElementById('input_gallery');1159 if (!gallery) return;1160 if (!gallery.matches(':hover')) return;1161 1162 const clipboardData = e.clipboardData || e.originalEvent.clipboardData;1163 if (!clipboardData) return;1164 1165 const items = clipboardData.items;1166 const files = [];1167 1168 for (let i = 0; i < items.length; i++) {1169 if (items[i].kind === 'file' && items[i].type.startsWith('image/')) {1170 files.push(items[i].getAsFile());1171 }1172 }1173 1174 if (files.length === 0 && clipboardData.files.length > 0) {1175 for (let i = 0; i < clipboardData.files.length; i++) {1176 if (clipboardData.files[i].type.startsWith('image/')) {1177 files.push(clipboardData.files[i]);1178 }1179 }1180 }1181 1182 if (files.length === 0) return;1183 1184 const uploadInput = gallery.querySelector('input[type="file"]');1185 if (uploadInput) {1186 e.preventDefault();1187 e.stopPropagation();1188 const dataTransfer = new DataTransfer();1189 files.forEach(file => dataTransfer.items.add(file));1190 uploadInput.files = dataTransfer.files;1191 uploadInput.dispatchEvent(new Event('change', { bubbles: true }));1192 }1193 });1194}1195"""1196 1197 1198# ----------------1199# Gradio layout1200# ----------------