"""Chain-of-Zoom runtime (https://github.com/bryanswkim/Chain-of-Zoom). optimized by DeepBeepMeep """ from __future__ import annotations import gc import math import os from dataclasses import dataclass import numpy as np import torch from PIL import Image from mmgp import offload from shared.utils import offload_registry COZ_LATENT_TILE_SIZE = 64 COZ_LATENT_TILE_OVERLAP = 16 COZ_STEP_FACTORS = {2.0: (2,), 4.0: (4,), 8.0: (4, 2), 16.0: (4, 4)} COZ_VLM_CONTEXT_SIZE = 512 COZ_VLM_MAX_NEW_TOKENS = 16 COZ_VLM_VISION_BATCH = 8 # tile image pairs encoded per vision tower pass COZ_VLM_PROMPT_BATCH = 16 # tile prompts decoded concurrently by the nano-vllm engine COZ_VLM_PROCESSOR_USE_FAST = True # fast (torchvision) image preprocessing; cheaper CPU prep, slightly different pixel values COZ_VLM_MESSAGE = "The second image is a zoom-in of the first image. Based on this knowledge, what is in the second image? Give me a set of English words." COZ_CLIP_MAX_LENGTH = 77 COZ_T5_MAX_LENGTH = 512 COZ_CLIP_TEXT_BATCH_SIZE = 1 COZ_T5_TEXT_BATCH_SIZE = 1 COZ_TIMESTEP = 1000.0 _CONFIGS_DIR = os.path.join(os.path.dirname(__file__), "configs") @dataclass class CoZPaths: transformer: str vae: str clip_l: str clip_tokenizer_dir: str # CLIP-L tokenizer files; CLIP-G uses the same vocab with pad token "!" clip_g: str t5: str vlm: str vlm_dir: str # folder holding the VLM config/processor/tokenizer files def _config_path(filename: str) -> str: return os.path.join(_CONFIGS_DIR, filename) def _flash_attention_2_available() -> bool: try: import flash_attn # noqa: F401 return True except Exception: return False def _abort_requested(abort_callback) -> bool: return callable(abort_callback) and abort_callback() def _report_progress(progress_callback, phase: str, current_step: int | None = None, total_steps: int | None = None, *, label_step: int | None = None, label_total: int | None = None) -> None: if callable(progress_callback): del label_step, label_total progress_callback(phase, current_step, total_steps) def _coz_phase(phase: str, step_label: str) -> str: return f"CoZ {phase} {step_label}" def _batch_count(total_items: int, batch_size: int) -> int: return (int(total_items) + int(batch_size) - 1) // int(batch_size) def _normalize_vlm_lm_engine(lm_decoder_engine: str | None) -> str: engine = str(lm_decoder_engine or "legacy").strip().lower() if engine == "cudagraph": engine = "cg" if engine not in ("legacy", "cg", "vllm"): raise ValueError(f"Unsupported Chain-of-Zoom VLM LM engine '{lm_decoder_engine}'. Expected one of: legacy, cg, vllm.") return engine def _disable_vlm_text_vllm_kernels(model) -> None: for module in model.modules(): if hasattr(module, "flash_attn_varlen_func"): module.flash_attn_varlen_func = None if hasattr(module, "flash_attn_with_kvcache"): module.flash_attn_with_kvcache = None if hasattr(module, "use_triton_kv_cache"): module.use_triton_kv_cache = False if hasattr(module, "use_triton_rmsnorm"): module.use_triton_rmsnorm = False def _grid_positions(length: int, tile: int, stride: int) -> list[int]: if length <= tile: return [0] positions = list(range(0, length - tile + 1, stride)) if positions[-1] != length - tile: positions.append(length - tile) return positions def _gaussian_weights(width: int, height: int, device) -> torch.Tensor: var = 0.01 midpoint_x = (width - 1) / 2 x_probs = [math.exp(-(x - midpoint_x) * (x - midpoint_x) / (width * width) / (2 * var)) / math.sqrt(2 * math.pi * var) for x in range(width)] midpoint_y = height / 2 y_probs = [math.exp(-(y - midpoint_y) * (y - midpoint_y) / (height * height) / (2 * var)) / math.sqrt(2 * math.pi * var) for y in range(height)] return torch.tensor(np.outer(y_probs, x_probs), device=device, dtype=torch.float32) class ChainOfZoomRuntime: def __init__(self) -> None: self.dtype = torch.bfloat16 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.transformer = None self.vae = None self.clip_l = None self.clip_g = None self.t5 = None self.vlm = None self.vlm_visual = None self.vlm_text = None self.vlm_engine = None self.vlm_config = None self.vlm_generation_config = None self.tokenizer_l = None self.tokenizer_g = None self.tokenizer_t5 = None self.vlm_processor = None self.vlm_llm = None self.offloadobj = None self.profile = None self.paths: CoZPaths | None = None self.use_vllm_decoder = False self.vlm_lm_engine = None self.vlm_vision_batch = COZ_VLM_VISION_BATCH self.vlm_prompt_batch = COZ_VLM_PROMPT_BATCH def load(self, paths: CoZPaths, profile, init_pipe, lm_decoder_engine: str = "legacy", vlm_vision_batch: int = COZ_VLM_VISION_BATCH, vlm_prompt_batch: int = COZ_VLM_PROMPT_BATCH) -> None: lm_decoder_engine = _normalize_vlm_lm_engine(lm_decoder_engine) vlm_vision_batch = int(vlm_vision_batch) vlm_prompt_batch = int(vlm_prompt_batch) use_vllm_decoder = lm_decoder_engine in ("cg", "vllm") if self.offloadobj is not None and self.paths == paths and self.profile == profile and self.vlm_lm_engine == lm_decoder_engine and self.vlm_vision_batch == vlm_vision_batch and self.vlm_prompt_batch == vlm_prompt_batch: return self.release() from accelerate import init_empty_weights from diffusers import AutoencoderKL from transformers import AutoConfig, AutoProcessor, CLIPTextModelWithProjection, CLIPTokenizer, GenerationConfig, Qwen2_5_VLForConditionalGeneration, T5Tokenizer from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel from shared.llm_engines.nanovllm.models.qwen2_5_vl import Qwen2_5_VLForCausalLM from shared.llm_engines.nanovllm.vllm_support import NanoVllmTextEngine from postprocessing.chain_of_zoom.sd3_transformer import SD3Transformer, SD3_MEDIUM_CONFIG with init_empty_weights(include_buffers=True): self.transformer = SD3Transformer(**SD3_MEDIUM_CONFIG) self.transformer.init_runtime_buffers() offload.load_model_data(self.transformer, paths.transformer, writable_tensors=False, default_dtype=self.dtype, ignore_unused_weights=True, verboseLevel=-1) self.vae = offload.fast_load_transformers_model(paths.vae, modelClass=AutoencoderKL, defaultConfigPath=_config_path("sd3_vae_config.json"), writable_tensors=False) self.vae.enable_tiling() self.clip_l = offload.fast_load_transformers_model(paths.clip_l, modelClass=CLIPTextModelWithProjection, forcedConfigPath=os.path.join(os.path.dirname(paths.clip_l), "config.json"), ignore_unused_weights=True, writable_tensors=False) self.clip_g = offload.fast_load_transformers_model(paths.clip_g, modelClass=CLIPTextModelWithProjection, forcedConfigPath=os.path.join(os.path.dirname(paths.clip_g), "config.json"), writable_tensors=False) self.t5 = offload.fast_load_transformers_model(paths.t5, writable_tensors=False) self.vlm_config = AutoConfig.from_pretrained(paths.vlm_dir) self.vlm_generation_config = GenerationConfig.from_pretrained(paths.vlm_dir) self.vlm_generation_config.do_sample = False self.vlm_generation_config.temperature = None self.use_vllm_decoder = use_vllm_decoder self.vlm_lm_engine = lm_decoder_engine self.vlm_vision_batch = vlm_vision_batch self.vlm_prompt_batch = vlm_prompt_batch print(f"[Chain-of-Zoom] VLM LM Engine='{self.vlm_lm_engine}'") print(f"[Chain-of-Zoom] VLM tile batch sizes: vision tower={self.vlm_vision_batch}, decoder={self.vlm_prompt_batch}") # Only vllm mode may use vLLM kernels. cg still uses nano-vLLM CUDA graphs, # but keeps the VLM on non-Flash2/non-Triton kernels. vision_attn = "flash_attention_2" if lm_decoder_engine == "vllm" and _flash_attention_2_available() else None if self.use_vllm_decoder: with init_empty_weights(include_buffers=True): self.vlm = Qwen2_5_VLForConditionalGeneration(self.vlm_config) self.vlm_text = Qwen2_5_VLForCausalLM(self.vlm_config.text_config) with init_empty_weights(include_buffers=False): if vision_attn: self.vlm_visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(self.vlm_config.vision_config, attn_implementation=vision_attn, torch_dtype=self.dtype) else: self.vlm_visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(self.vlm_config.vision_config) offload.load_model_data(self.vlm_visual, paths.vlm, modelPrefix="model.visual", writable_tensors=False, default_dtype=self.dtype, ignore_unused_weights=True, verboseLevel=-1) offload.load_model_data(self.vlm_text, paths.vlm, modelPrefix="model.language_model", writable_tensors=False, default_dtype=self.dtype, ignore_unused_weights=True, verboseLevel=-1) if lm_decoder_engine == "cg": _disable_vlm_text_vllm_kernels(self.vlm_text) else: self.vlm = offload.fast_load_transformers_model(paths.vlm, modelClass=Qwen2_5_VLForConditionalGeneration, defaultConfigPath=os.path.join(paths.vlm_dir, "config.json"), writable_tensors=False) self.vlm.tie_weights() self.vlm.generation_config = self.vlm_generation_config if vision_attn: self.vlm.model.visual.config._attn_implementation = vision_attn for model in (self.transformer, self.vae, self.clip_l, self.clip_g, self.t5, self.vlm_visual if self.use_vllm_decoder else self.vlm, self.vlm_text): if model is None: continue model.eval().requires_grad_(False) self.tokenizer_l = CLIPTokenizer.from_pretrained(paths.clip_tokenizer_dir) # SD3's tokenizer_2 is the same CLIP vocab with pad token "!" (id 0) self.tokenizer_g = CLIPTokenizer.from_pretrained(paths.clip_tokenizer_dir, pad_token="!") self.tokenizer_t5 = T5Tokenizer.from_pretrained(os.path.dirname(paths.t5)) self.vlm_processor = AutoProcessor.from_pretrained(paths.vlm_dir, use_fast=COZ_VLM_PROCESSOR_USE_FAST) if self.use_vllm_decoder: self.vlm_engine = NanoVllmTextEngine(model=self.vlm_text, model_path=paths.vlm_dir, tokenizer=self.vlm_processor.tokenizer, enforce_eager=False) self.vlm_llm = self.vlm_text else: self.vlm_llm = torch.nn.ModuleDict({"language_model": self.vlm.model.language_model, "lm_head": self.vlm.lm_head}) pipe = { "transformer": self.transformer, "vae": self.vae, "clip_l": self.clip_l, "clip_g": self.clip_g, "t5": self.t5, "vlm_visual": self.vlm_visual if self.use_vllm_decoder else self.vlm.model.visual, "vlm_llm": self.vlm_llm, } kwargs = {} profile_no = init_pipe(pipe, kwargs, profile) kwargs.setdefault("budgets", {})["clip_l"] = 0 kwargs["budgets"]["clip_g"] = 0 kwargs["budgets"]["vlm_visual"] = 0 if self.use_vllm_decoder: kwargs["budgets"]["vlm_llm"] = 0 kwargs.setdefault("coTenantsMap", {}).update({"clip_l": ["clip_g"], "clip_g": ["clip_l"]}) kwargs["pinnedMemory"] = False self.offloadobj = offload.profile(pipe, profile_no=profile_no, quantizeTransformer=False, convertWeightsFloatTo=self.dtype, verboseLevel=-1, **kwargs) offload_registry.register_offloadobj("Chain-of-Zoom", self.offloadobj, self.release) self.profile = profile self.paths = paths def _unload_mmgp(self) -> None: if self.vlm_engine is not None: self.vlm_engine.release_runtime_allocations() if self.offloadobj is not None: self.offloadobj.unload_all() if torch.cuda.is_available(): torch.cuda.empty_cache() def _current_vlm_engine_label(self) -> str | None: if self.vlm_lm_engine is None: return None return f"LM Engine='{self.vlm_lm_engine}'" def release(self) -> None: vlm_engine_label = self._current_vlm_engine_label() if vlm_engine_label is not None: print(f"[Chain-of-Zoom] Closing VLM {vlm_engine_label}") # CUDA graphs hold model/KV pointers; clear them before MMGP releases model storage. if self.vlm_engine is not None: self.vlm_engine.close() self.vlm_engine = None if self.offloadobj is not None: offload_registry.unregister_offloadobj("Chain-of-Zoom", self.offloadobj) self.offloadobj.release() self.offloadobj = None self.transformer = None self.vae = None self.clip_l = None self.clip_g = None self.t5 = None self.vlm = None self.vlm_visual = None self.vlm_text = None self.vlm_config = None self.vlm_generation_config = None self.tokenizer_l = None self.tokenizer_g = None self.tokenizer_t5 = None self.vlm_processor = None self.vlm_llm = None self.use_vllm_decoder = False self.vlm_lm_engine = None self.vlm_vision_batch = COZ_VLM_VISION_BATCH self.vlm_prompt_batch = COZ_VLM_PROMPT_BATCH self.profile = None self.paths = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def _vlm_inputs(self, context_image: Image.Image, tile_image: Image.Image): messages = [ {"role": "system", "content": COZ_VLM_MESSAGE}, {"role": "user", "content": [{"type": "image"}, {"type": "image"}]}, ] text = self.vlm_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) return self.vlm_processor(text=[text], images=[context_image, tile_image], padding=True, return_tensors="pt") def _vlm_visual_embeds(self, inputs_batch: list) -> list[torch.Tensor]: """One vision tower pass over the images of several tile jobs; returns per-job packed embeds.""" pixel_values = torch.cat([inputs["pixel_values"] for inputs in inputs_batch], dim=0).to(self.device) image_grid_thw = torch.cat([inputs["image_grid_thw"] for inputs in inputs_batch], dim=0).to(self.device) if self.use_vllm_decoder: pixel_values = pixel_values.type(self.vlm_visual.dtype) image_embeds = self.vlm_visual(pixel_values, grid_thw=image_grid_thw) merge_size = self.vlm_visual.spatial_merge_size else: image_embeds = torch.cat(self.vlm.model.get_image_features(pixel_values, image_grid_thw), dim=0) merge_size = self.vlm.model.visual.spatial_merge_size tokens_per_image = (image_grid_thw.prod(-1) // merge_size**2).tolist() job_sizes = [] image_index = 0 for inputs in inputs_batch: image_count = int(inputs["image_grid_thw"].shape[0]) job_sizes.append(int(sum(tokens_per_image[image_index:image_index + image_count]))) image_index += image_count image_embeds = image_embeds.to(self.dtype).cpu() del pixel_values, image_grid_thw return list(torch.split(image_embeds, job_sizes, dim=0)) def _vlm_prompts_from_embeds_vllm(self, jobs: list[dict]) -> list[str]: """Decode the prompts of several tile jobs concurrently through the nano-vllm engine.""" requests = [] for job in jobs: input_ids = job["input_ids"].to(self.device) attention_mask = job.get("attention_mask") attention_mask = attention_mask.to(self.device) if attention_mask is not None else torch.ones_like(input_ids, device=self.device) image_grid_thw = job["image_grid_thw"].to(self.device) image_embeds = job["image_embeds"] inputs_embeds = self.vlm_text.embed_tokens(input_ids) image_mask = input_ids == self.vlm_config.image_token_id if int(image_mask.sum()) != image_embeds.shape[0]: raise ValueError(f"VLM image token/features mismatch: tokens={int(image_mask.sum())}, features={image_embeds.shape[0]}") image_embeds = image_embeds.to(device=self.device, dtype=inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_mask.unsqueeze(-1).expand_as(inputs_embeds), image_embeds) position_ids, rope_deltas = self.vlm.model.get_rope_index(input_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask) active_mask = attention_mask[0].bool() requests.append({ "prompt_token_ids": [int(token_id) for token_id in input_ids[0][active_mask].tolist()], "prompt_embeds": inputs_embeds[0, active_mask].contiguous().cpu(), "prompt_position_ids": position_ids[:, 0, active_mask].contiguous().cpu(), "position_offset": int(rope_deltas.reshape(-1)[0].item()) if torch.is_tensor(rope_deltas) and rope_deltas.numel() > 0 else 0, }) del input_ids, attention_mask, image_grid_thw, inputs_embeds, image_embeds, position_ids, rope_deltas results = self.vlm_engine.generate_embedded_batch( requests, max_tokens=COZ_VLM_MAX_NEW_TOKENS, temperature=None, top_p=None, top_k=1, cfg_scale=1.0, seed=None, use_tqdm=False, release_vram_after=False, ignore_eos=False, repetition_penalty=float(self.vlm_generation_config.repetition_penalty or 1.0), ) requests.clear() token_id_lists = [[] if result is None else result.get("token_ids", []) for result in results or []] return [prompt.strip() for prompt in self.vlm_processor.batch_decode(token_id_lists, skip_special_tokens=True, clean_up_tokenization_spaces=False)] def _vlm_prompts_from_embeds_legacy(self, jobs: list[dict]) -> list[str]: input_ids = torch.cat([job["input_ids"] for job in jobs], dim=0).to(self.device) attention_mask = torch.cat([ job["attention_mask"] if job.get("attention_mask") is not None else torch.ones_like(job["input_ids"]) for job in jobs ], dim=0).to(self.device) image_grid_thw = torch.cat([job["image_grid_thw"] for job in jobs], dim=0).to(self.device) image_embeds = torch.cat([job["image_embeds"] for job in jobs], dim=0) inputs_embeds = self.vlm.model.get_input_embeddings()(input_ids) image_mask = input_ids == self.vlm.config.image_token_id if int(image_mask.sum()) != image_embeds.shape[0]: raise ValueError(f"VLM image token/features mismatch: tokens={int(image_mask.sum())}, features={image_embeds.shape[0]}") image_embeds = image_embeds.to(device=self.device, dtype=inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_mask.unsqueeze(-1).expand_as(inputs_embeds), image_embeds) self.vlm.model.rope_deltas = None position_ids, rope_deltas = self.vlm.model.get_rope_index(input_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask) self.vlm.model.rope_deltas = rope_deltas eos_token_ids = self.vlm_generation_config.eos_token_id eos_token_ids = {eos_token_ids} if isinstance(eos_token_ids, int) else set(eos_token_ids or []) forced_eos_token_id = int(next(iter(eos_token_ids))) if eos_token_ids else None repetition_penalty = float(self.vlm_generation_config.repetition_penalty or 1.0) generated = [] finished = torch.zeros((input_ids.shape[0], 1), device=self.device, dtype=torch.bool) past_key_values = None current_input_ids = input_ids current_inputs_embeds = inputs_embeds current_position_ids = position_ids cache_position = torch.arange(input_ids.shape[1], device=self.device, dtype=torch.long) for _ in range(COZ_VLM_MAX_NEW_TOKENS): output = self.vlm( input_ids=current_input_ids, inputs_embeds=current_inputs_embeds, attention_mask=attention_mask, position_ids=current_position_ids, image_grid_thw=image_grid_thw, past_key_values=past_key_values, use_cache=True, cache_position=cache_position, logits_to_keep=1, ) logits = output.logits[:, -1, :] if repetition_penalty != 1.0: previous_tokens = torch.cat([input_ids] + generated, dim=1) if generated else input_ids for row_no in range(input_ids.shape[0]): previous_ids = torch.unique(previous_tokens[row_no]) previous_scores = logits[row_no, previous_ids] logits[row_no, previous_ids] = torch.where(previous_scores < 0, previous_scores * repetition_penalty, previous_scores / repetition_penalty) next_token = logits.argmax(dim=-1, keepdim=True) if forced_eos_token_id is not None: next_token = torch.where(finished, torch.full_like(next_token, forced_eos_token_id), next_token) generated.append(next_token) past_key_values = output.past_key_values del output, logits if eos_token_ids: token_finished = torch.zeros_like(finished) for eos_token_id in eos_token_ids: token_finished |= next_token == int(eos_token_id) finished |= token_finished if bool(finished.all()): break attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=1) cache_position = cache_position[-1:] + 1 current_input_ids = next_token current_inputs_embeds = None current_position_ids = cache_position.view(1, 1, -1).expand(3, input_ids.shape[0], -1) + rope_deltas.view(1, input_ids.shape[0], 1).to(self.device) generated_ids = torch.cat(generated, dim=1) prompts = self.vlm_processor.batch_decode(generated_ids.cpu(), skip_special_tokens=True, clean_up_tokenization_spaces=False) del input_ids, attention_mask, image_grid_thw, inputs_embeds, image_embeds, position_ids, rope_deltas, generated_ids, generated, finished, past_key_values return [prompt.strip() for prompt in prompts] def _encode_clip_prompts(self, text_encoder, tokenizer, prompts: list[str], phase: str, progress_start: int, progress_total: int, abort_callback=None, progress_callback=None) -> tuple[torch.Tensor, torch.Tensor] | None: hidden_batches = [] pooled_batches = [] total = len(prompts) for batch_no, start in enumerate(range(0, total, COZ_CLIP_TEXT_BATCH_SIZE)): if _abort_requested(abort_callback): return None current_step = progress_start + batch_no + 1 _report_progress(progress_callback, phase, current_step, progress_total) ids = tokenizer(prompts[start:start + COZ_CLIP_TEXT_BATCH_SIZE], padding="max_length", max_length=COZ_CLIP_MAX_LENGTH, truncation=True, return_tensors="pt").input_ids output = text_encoder(ids.to(self.device), output_hidden_states=True) hidden_batches.append(output.hidden_states[-2].to(self.dtype).cpu()) pooled_batches.append(output[0].to(self.dtype).cpu()) del ids, output return torch.cat(hidden_batches), torch.cat(pooled_batches) def _encode_t5_prompts(self, prompts: list[str], phase: str, progress_start: int, progress_total: int, abort_callback=None, progress_callback=None) -> torch.Tensor | None: batches = [] total = len(prompts) for batch_no, start in enumerate(range(0, total, COZ_T5_TEXT_BATCH_SIZE)): if _abort_requested(abort_callback): return None current_step = progress_start + batch_no + 1 _report_progress(progress_callback, phase, current_step, progress_total) ids = self.tokenizer_t5(prompts[start:start + COZ_T5_TEXT_BATCH_SIZE], padding="max_length", max_length=COZ_T5_MAX_LENGTH, truncation=True, add_special_tokens=True, return_tensors="pt").input_ids batches.append(self.t5(ids.to(self.device))[0].to(self.dtype).cpu()) del ids return torch.cat(batches) def _encode_prompts(self, prompts: list[str], step_label: str, abort_callback=None, progress_callback=None) -> dict[str, tuple[torch.Tensor, torch.Tensor]] | None: unique_prompts = list(dict.fromkeys(prompts)) clip_batches = _batch_count(len(unique_prompts), COZ_CLIP_TEXT_BATCH_SIZE) t5_batches = _batch_count(len(unique_prompts), COZ_T5_TEXT_BATCH_SIZE) text_progress_total = clip_batches * 2 + t5_batches clip1 = self._encode_clip_prompts(self.clip_l, self.tokenizer_l, unique_prompts, _coz_phase("Text Encoding CLIP-L", step_label), 0, text_progress_total, abort_callback=abort_callback, progress_callback=progress_callback) if clip1 is None: return None clip2 = self._encode_clip_prompts(self.clip_g, self.tokenizer_g, unique_prompts, _coz_phase("Text Encoding CLIP-G", step_label), clip_batches, text_progress_total, abort_callback=abort_callback, progress_callback=progress_callback) if clip2 is None: return None t5_emb = self._encode_t5_prompts(unique_prompts, _coz_phase("Text Encoding T5", step_label), clip_batches * 2, text_progress_total, abort_callback=abort_callback, progress_callback=progress_callback) if t5_emb is None: return None embeds_cache = {} clip1_emb, pool1 = clip1 clip2_emb, pool2 = clip2 for prompt_no, prompt in enumerate(unique_prompts): clip_emb = torch.cat([clip1_emb[prompt_no:prompt_no + 1], clip2_emb[prompt_no:prompt_no + 1]], dim=-1) clip_emb = torch.nn.functional.pad(clip_emb, (0, t5_emb.shape[-1] - clip_emb.shape[-1])) prompt_emb = torch.cat([clip_emb, t5_emb[prompt_no:prompt_no + 1]], dim=-2) pooled_emb = torch.cat([pool1[prompt_no:prompt_no + 1], pool2[prompt_no:prompt_no + 1]], dim=-1) embeds_cache[prompt] = (prompt_emb, pooled_emb) return embeds_cache def _sr_step(self, image: Image.Image, context_image: Image.Image, out_width: int, out_height: int, seed: int, step_label: str, abort_callback=None, progress_callback=None) -> Image.Image | None: up_image = image.resize((out_width, out_height), Image.LANCZOS) latent_height, latent_width = out_height // 8, out_width // 8 tile_h = min(COZ_LATENT_TILE_SIZE, latent_height) tile_w = min(COZ_LATENT_TILE_SIZE, latent_width) stride_h = max(1, tile_h - COZ_LATENT_TILE_OVERLAP) stride_w = max(1, tile_w - COZ_LATENT_TILE_OVERLAP) positions = [(y0, x0) for y0 in _grid_positions(latent_height, tile_h, stride_h) for x0 in _grid_positions(latent_width, tile_w, stride_w)] total_tiles = len(positions) # phase 1: multi-scale aware prompt per tile (original image as zoom-out context, # tile crop of the pre-upsampled image as zoom-in; equivalent to upstream's decoded # latent patch, without the extra per-tile VAE decode). Split Qwen2.5-VL into a # visual pass over every tile, then an LLM pass over cached visual embeddings. # every tile shares the same chat template and image sizes, so the prompt token ids # and the context image preprocessing are computed once; per tile only the crop # itself goes through the (CPU bound) image processor, on a small thread pool from concurrent.futures import ThreadPoolExecutor from shared.utils.utils import get_default_workers first_inputs = self._vlm_inputs(context_image, up_image.crop((positions[0][1] * 8, positions[0][0] * 8, (positions[0][1] + tile_w) * 8, (positions[0][0] + tile_h) * 8))) # the fast image processor may return CUDA tensors; jobs are staged on CPU first_inputs = {key: value.cpu() if torch.is_tensor(value) else value for key, value in first_inputs.items()} context_patches = int(first_inputs["image_grid_thw"][0].prod()) template = { "input_ids": first_inputs["input_ids"].cpu(), "attention_mask": first_inputs["attention_mask"].cpu() if "attention_mask" in first_inputs else None, "context_grid": first_inputs["image_grid_thw"][:1].cpu(), "context_pixel_values": first_inputs["pixel_values"][:context_patches].cpu(), "tile_grid": first_inputs["image_grid_thw"][1:].cpu(), } def tile_inputs(tile_image: Image.Image) -> dict: tile_features = self.vlm_processor.image_processor(images=[tile_image], return_tensors="pt") tile_grid = tile_features["image_grid_thw"].cpu() if not torch.equal(tile_grid, template["tile_grid"]): raise RuntimeError(f"VLM tile grid mismatch: {tile_grid.tolist()} vs {template['tile_grid'].tolist()}") return { "input_ids": template["input_ids"], "attention_mask": template["attention_mask"], "image_grid_thw": torch.cat([template["context_grid"], tile_grid], dim=0), "pixel_values": torch.cat([template["context_pixel_values"], tile_features["pixel_values"].cpu()], dim=0), } vlm_jobs = [] vision_phase = _coz_phase("VLM Vision", step_label) with ThreadPoolExecutor(max_workers=max(1, min(int(get_default_workers()), self.vlm_vision_batch))) as executor: for batch_start in range(0, total_tiles, self.vlm_vision_batch): if _abort_requested(abort_callback): return None batch_positions = positions[batch_start:batch_start + self.vlm_vision_batch] _report_progress(progress_callback, vision_phase, batch_start + len(batch_positions), total_tiles) if batch_start == 0: batch_positions = batch_positions[1:] tile_images = [up_image.crop((x0 * 8, y0 * 8, (x0 + tile_w) * 8, (y0 + tile_h) * 8)) for y0, x0 in batch_positions] inputs_batch = ([first_inputs] if batch_start == 0 else []) + list(executor.map(tile_inputs, tile_images)) del tile_images image_embeds_batch = self._vlm_visual_embeds(inputs_batch) for inputs, image_embeds in zip(inputs_batch, image_embeds_batch): attention_mask = inputs.get("attention_mask") vlm_jobs.append({ "input_ids": inputs["input_ids"].cpu(), "attention_mask": attention_mask.cpu() if attention_mask is not None else None, "image_grid_thw": inputs["image_grid_thw"].cpu(), "image_embeds": image_embeds, }) del inputs_batch, image_embeds_batch first_inputs = template = None prompts = [] text_phase = _coz_phase("VLM Text", step_label) if self.use_vllm_decoder and vlm_jobs: max_prompt_len = max(int(job["attention_mask"].sum().item()) if job["attention_mask"] is not None else int(job["input_ids"].shape[-1]) for job in vlm_jobs) self.vlm_engine.reserve_runtime(prompt_len=max_prompt_len, max_tokens=COZ_VLM_MAX_NEW_TOKENS, cfg_scale=1.0, num_seqs=min(self.vlm_prompt_batch, len(vlm_jobs))) for batch_start in range(0, len(vlm_jobs), self.vlm_prompt_batch): if _abort_requested(abort_callback): return None batch_jobs = vlm_jobs[batch_start:batch_start + self.vlm_prompt_batch] _report_progress(progress_callback, text_phase, batch_start + len(batch_jobs), total_tiles) for tile_offset, prompt in enumerate(self._vlm_prompts_from_embeds_vllm(batch_jobs)): prompts.append(prompt) print(f"[Chain-of-Zoom] {step_label} VLM tile {batch_start + tile_offset + 1}/{total_tiles} prompt: {prompt}") for job in batch_jobs: job.clear() else: for batch_start in range(0, len(vlm_jobs), self.vlm_prompt_batch): if _abort_requested(abort_callback): return None batch_jobs = vlm_jobs[batch_start:batch_start + self.vlm_prompt_batch] _report_progress(progress_callback, text_phase, batch_start + len(batch_jobs), total_tiles) for tile_offset, prompt in enumerate(self._vlm_prompts_from_embeds_legacy(batch_jobs)): prompts.append(prompt) print(f"[Chain-of-Zoom] {step_label} VLM tile {batch_start + tile_offset + 1}/{total_tiles} prompt: {prompt}") for job in batch_jobs: job.clear() vlm_jobs.clear() if self.use_vllm_decoder and self.vlm_engine is not None: # VLM phase over: free the engine KV cache + CUDA graph pool before the # text-encode/VAE/diffusion phases. MMGP will move the weights anyway, which # invalidates the captured graphs; the engine rebuilds them next VLM phase. self.vlm_engine.release_runtime_allocations() # phase 2: text embeddings (cached per unique prompt, stored on CPU) embeds_cache = self._encode_prompts(prompts, step_label, abort_callback=abort_callback, progress_callback=progress_callback) if embeds_cache is None: return None # phase 3: VAE encode (OSEDiff convention: scaling factor only, no shift) if _abort_requested(abort_callback): return None _report_progress(progress_callback, _coz_phase("VAE Encode", step_label)) x_full = torch.from_numpy(np.array(up_image)).permute(2, 0, 1).unsqueeze(0).to(device=self.device, dtype=self.dtype).div_(127.5).sub_(1.0) generator = torch.Generator(device=self.device).manual_seed(seed) z_full = self.vae.encode(x_full).latent_dist.sample(generator) * self.vae.config.scaling_factor del x_full # phase 4: one-step OSEDiff on each tile, Gaussian-blended accumulation weights = _gaussian_weights(tile_w, tile_h, self.device) z_acc = torch.zeros_like(z_full, dtype=torch.float32) norm = torch.zeros_like(z_full, dtype=torch.float32) timestep = torch.tensor([COZ_TIMESTEP], device=self.device, dtype=self.dtype) diffusion_phase = _coz_phase("Diffusion", step_label) for tile_no, (y0, x0) in enumerate(positions): if _abort_requested(abort_callback): return None _report_progress(progress_callback, diffusion_phase, tile_no + 1, total_tiles) prompt_emb, pooled_emb = embeds_cache[prompts[tile_no]] patch = z_full[:, :, y0:y0 + tile_h, x0:x0 + tile_w] pred_v = self.transformer(patch, timestep, prompt_emb.to(device=self.device, dtype=self.dtype), pooled_emb.to(device=self.device, dtype=self.dtype)) tile_out = (patch - pred_v).float() z_acc[:, :, y0:y0 + tile_h, x0:x0 + tile_w] += tile_out * weights norm[:, :, y0:y0 + tile_h, x0:x0 + tile_w] += weights del patch, pred_v, tile_out embeds_cache.clear() z_full = (z_acc / (norm + 1e-10)).to(self.dtype) del z_acc, norm # phase 5: VAE decode if _abort_requested(abort_callback): return None _report_progress(progress_callback, _coz_phase("VAE Decode", step_label)) output = self.vae.decode(z_full / self.vae.config.scaling_factor, return_dict=False)[0].clamp(-1, 1) del z_full output = output[0].float().add_(1.0).mul_(127.5).round_().clamp_(0, 255).to(torch.uint8).permute(1, 2, 0).cpu().numpy() return Image.fromarray(output) @torch.inference_mode() def upscale(self, sample: torch.Tensor, scale: float, *, seed: int = 0, abort_callback=None, progress_callback=None) -> torch.Tensor | None: if self.device.type != "cuda": raise RuntimeError("Chain-of-Zoom requires CUDA.") factors = COZ_STEP_FACTORS[float(scale)] frame = sample[:, 0] if frame.dtype == torch.uint8: frame_np = frame.permute(1, 2, 0).cpu().numpy() else: frame_np = frame.permute(1, 2, 0).float().add(1.0).mul(127.5).round().clamp(0, 255).to(torch.uint8).cpu().numpy() if frame_np.shape[2] == 1: frame_np = np.repeat(frame_np, 3, axis=2) image = Image.fromarray(frame_np) target_w = int(round(image.width * scale / 16) * 16) target_h = int(round(image.height * scale / 16) * 16) # constant zoom-out context across the chain: the original input resized to 512 min side context_scale = COZ_VLM_CONTEXT_SIZE / min(image.width, image.height) context_image = image.resize((max(1, int(image.width * context_scale)), max(1, int(image.height * context_scale))), Image.LANCZOS) print(f"[Chain-of-Zoom] x{scale:g} upsampling in {len(factors)} step(s): input={image.width}x{image.height}, output={target_w}x{target_h}") for step_no, factor in enumerate(factors): if step_no == len(factors) - 1: step_w, step_h = target_w, target_h else: step_w = int(round(image.width * factor / 16) * 16) step_h = int(round(image.height * factor / 16) * 16) image = self._sr_step(image, context_image, step_w, step_h, seed, f"Step {step_no + 1}/{len(factors)}", abort_callback=abort_callback, progress_callback=progress_callback) if image is None: return None output = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(1) return output _RUNTIME = ChainOfZoomRuntime() def load_models(paths: CoZPaths, *, init_pipe, profile, lm_decoder_engine: str = "legacy", vlm_vision_batch: int = COZ_VLM_VISION_BATCH, vlm_prompt_batch: int = COZ_VLM_PROMPT_BATCH, progress_callback=None) -> None: _report_progress(progress_callback, "CoZ Loading Models") _RUNTIME.load(paths, profile=profile, init_pipe=init_pipe, lm_decoder_engine=lm_decoder_engine, vlm_vision_batch=vlm_vision_batch, vlm_prompt_batch=vlm_prompt_batch) def upscale_image(sample: torch.Tensor, scale: float, *, seed: int = 0, abort_callback=None, progress_callback=None) -> torch.Tensor | None: return _RUNTIME.upscale(sample, scale, seed=seed, abort_callback=abort_callback, progress_callback=progress_callback) def release_models() -> None: _RUNTIME.release()