Spaces:
Running on Zero
Running on Zero
File size: 5,880 Bytes
ee69ab7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | from __future__ import annotations
import os
import sys
from pathlib import Path
from .config import InferenceConfig
_BANNER = """
███████╗██╗ █████╗ ███████╗██╗ ██╗██╗ ██╗███████╗██████╗ ██╗
██╔════╝██║ ██╔══██╗██╔════╝██║ ██║██║ ██║██╔════╝██╔══██╗ ██║
█████╗ ██║ ███████║███████╗███████║██║ ██║███████╗██████╔╝ ████████╗
██╔══╝ ██║ ██╔══██║╚════██║██╔══██║╚██╗ ██╔╝╚════██║██╔══██╗ ╚══██╔══╝
██║ ███████╗██║ ██║███████║██║ ██║ ╚████╔╝ ███████║██║ ██║ ██║
╚═╝ ╚══════╝╚═╝ ╚═╝╚══════╝╚═╝ ╚═╝ ╚═══╝ ╚══════╝╚═╝ ╚═╝ ╚═╝
""".strip("\n")
def _supports_color() -> bool:
return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None
def _paint(text: str, code: str) -> str:
return f"\033[{code}m{text}\033[0m" if _supports_color() else text
def _path_label(value: str | None) -> str:
if not value:
return "Not specified"
path = Path(value)
parts = path.parts
if len(parts) <= 4:
return str(path)
return str(Path("…", *parts[-3:]))
def print_banner() -> None:
gradient = ("1;36", "1;36", "1;34", "1;34", "1;35", "1;35")
for line, color in zip(_BANNER.splitlines(), gradient):
print(_paint(line, color))
print(_paint(" STRUCTURED VIDEO SUPER-RESOLUTION", "1;37"))
print()
def print_configuration(config: InferenceConfig, args) -> None:
mode = f"STREAM · {config.streaming.chunk_latents}-latent chunks"
attention = config.attention_config()
decoder = "Tiny Decoder" if config.decoder.kind == "tiny" else "VAE Decoder"
sliding_range = (
config.sliding_local_range
if config.sliding_spatial_mode == "local"
else config.sparse_local_range
)
rows = (
("Mode", mode),
("Attention", f"{attention.backend} / {attention.kernel}"),
(
"Sparse Top-k",
str(config.sparse_topk or config.sparse_topk_ratio)
if config.is_sparse
else "Disabled",
),
("Local Range", str(config.sparse_local_range or "Disabled")),
(
"Hybrid Sliding",
(
"range "
f"{sliding_range} / "
f"spatial {config.sliding_spatial_mode} / "
f"temporal {config.sliding_temporal_mode} / "
f"cache {config.sliding_cache_temporal_windows} window(s) / "
f"{config.sliding_temporal_block_size} temporal windows"
if config.is_hybrid_sliding
else "Disabled"
),
),
("KV Time Groups", str(int(config.sparse_kv_ratio))),
(
"KV Cache Policy",
config.stream_cache_offload_mode,
),
("Decoder", decoder),
(
"Tiny Temporal Mode",
(
(f"{config.decoder.tiny_decoder_parallel_mode} · adaptive"
if config.decoder.tiny_decoder_parallel_mode == "auto"
else config.decoder.tiny_decoder_parallel_mode)
+ (
" · channels-last"
if config.decoder.tiny_decoder_channels_last
else " · contiguous"
)
)
if config.decoder.kind == "tiny"
else "N/A",
),
("Temporal Compression", f"{config.temporal_compression}×"),
(
"Color Correction",
(
f"Wavelet · {config.color_correction.levels} levels"
if config.color_correction.enabled
else "Disabled"
),
),
(
"Temporal Merge",
config.temporal_merge if config.temporal_compression < 4 else "N/A",
),
("Precision / Device", f"{config.dtype} / {config.device}"),
("Output Size", f"{args.width or 'Original'} × {args.height or 'Original'}"),
("Maximum Frames", str(args.max_frames or "All")),
("Steps / Seed", f"{config.num_inference_steps} / {config.seed}"),
("Sigma Shift", str(config.sigma_shift)),
)
print(_paint("┌─ Runtime Configuration", "1;35"))
width = max(len(label) for label, _ in rows)
for label, value in rows:
print(f"│ {label:<{width}} {_paint(value, '1;37')}")
print(_paint("└────────────────────────────────────────────────────", "2;35"))
print()
components = (
("DiT Model", args.dit),
("LQ Projection", args.lq_projection),
("Text Context", args.context),
(
"Tiny Decoder" if config.decoder.kind == "tiny" else "VAE Decoder",
args.tiny_decoder if config.decoder.kind == "tiny" else args.vae,
),
)
print(_paint("Model Components", "1;35"))
for name, path in components:
print(f" {_paint('◆', '36')} {name:<16} {_path_label(path)}")
print()
def print_stage(message: str, *, done: bool = False) -> None:
marker = _paint("✓", "1;32") if done else _paint("●", "1;36")
print(f"{marker} {message}", flush=True)
|