Add self-contained inference bundle (vendored CoZ code + ckpts + runner)

ckpt/SR_LoRA/._model_20001.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e9aed9ddeffb56d172d7aede994956b3b2c425d0f3943ca359ccd454d71727c
+size 163

ckpt/SR_LoRA/model_20001.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:697d3f9ab69a222006ca3ae48503cf057774c8142646301e9bba90e58242e47e
+size 8111108

ckpt/SR_VAE/._vae_encoder_20001.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e9aed9ddeffb56d172d7aede994956b3b2c425d0f3943ca359ccd454d71727c
+size 163

ckpt/SR_VAE/vae_encoder_20001.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed7f7aa03dfcbce9016d51c5aa8d3920428b3d7c9a678c721cd062d01805ae4a
+size 69346330

ckpt/VLM_LoRA/checkpoint-10000/._adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e9aed9ddeffb56d172d7aede994956b3b2c425d0f3943ca359ccd454d71727c
+size 163

ckpt/VLM_LoRA/checkpoint-10000/adapter_config.json

@@ -0,0 +1,31 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "/mnt/data1/bryanswkim/cache/modelscope/models/Qwen/Qwen2___5-VL-3B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": [],
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": "^(model).*\\.(q_proj|o_proj|gate_proj|down_proj|k_proj|up_proj|v_proj)$",
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_rslora": false
+}

ckpt/VLM_LoRA/checkpoint-10000/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ba6fc24e76e0e078ceb6e067c49bcfbe86cb0d8995add1515172d922fba2ebd
+size 59933632

coz/lora/lora_layers.py

@@ -0,0 +1,137 @@
+from typing import List, Optional, Set, Type, Union
+
+import torch
+from torch import nn
+
+
+class LoraInjectedLinear(nn.Module):
+    """
+    Linear layer with LoRA injection.
+    Taken from https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+    """
+    def __init__(
+        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
+    ):
+        super().__init__()
+
+        if r > min(in_features, out_features):
+            raise ValueError(
+                f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
+            )
+        self.r = r
+        self.linear = nn.Linear(in_features, out_features, bias)
+        self.lora_down = nn.Linear(in_features, r, bias=False)
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Linear(r, out_features, bias=False)
+        self.scale = scale
+        self.selector = nn.Identity()
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.linear(input.float())
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input.float()))))
+            * self.scale
+        ).half()
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Linear(self.r, self.r, bias=False)
+        self.selector.weight.data = torch.diag(diag)
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+class LoraInjectedConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups: int = 1,
+        bias: bool = True,
+        r: int = 4,
+        dropout_p: float = 0.1,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        if r > min(in_channels, out_channels):
+            raise ValueError(
+                f"LoRA rank {r} must be less or equal than {min(in_channels, out_channels)}"
+            )
+        self.r = r
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+        self.lora_down = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=r,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=False,
+        )
+        self.dropout = nn.Dropout(dropout_p)
+        self.lora_up = nn.Conv2d(
+            in_channels=r,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector = nn.Identity()
+        self.scale = scale
+
+        nn.init.normal_(self.lora_down.weight, std=1 / r)
+        nn.init.zeros_(self.lora_up.weight)
+
+    def forward(self, input):
+        return (
+            self.conv(input)
+            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
+            * self.scale
+        )
+
+    def realize_as_lora(self):
+        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
+
+    def set_selector_from_diag(self, diag: torch.Tensor):
+        # diag is a 1D tensor of size (r,)
+        assert diag.shape == (self.r,)
+        self.selector = nn.Conv2d(
+            in_channels=self.r,
+            out_channels=self.r,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.selector.weight.data = torch.diag(diag)
+
+        # same device + dtype as lora_up
+        self.selector.weight.data = self.selector.weight.data.to(
+            self.lora_up.weight.device
+        ).to(self.lora_up.weight.dtype)
+
+
+

coz/lora/lora_utils.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+from lora.lora_layers import LoraInjectedLinear, LoraInjectedConv2d
+
+def _find_modules(model, ancestor_class=None, search_class=[nn.Linear], exclude_children_of=[LoraInjectedLinear]):
+    # Get the targets we should replace all linears under
+    if ancestor_class is not None:
+        ancestors = (
+            module
+            for module in model.modules()
+            if module.__class__.__name__ in ancestor_class
+        )
+    else:
+        # this, incase you want to naively iterate over all modules.
+        ancestors = [module for module in model.modules()]
+
+    for ancestor in ancestors:
+        for fullname, module in ancestor.named_modules():
+            # if 'norm1_context' in fullname:
+                if any([isinstance(module, _class) for _class in search_class]):
+                    *path, name = fullname.split(".")
+                    parent = ancestor
+                    while path:
+                        parent = parent.get_submodule(path.pop(0))
+                    if exclude_children_of and any(
+                        [isinstance(parent, _class) for _class in exclude_children_of]
+                    ):
+                        continue
+                    yield parent, name, module
+
+def extract_lora_ups_down(model, target_replace_module={'AdaLayerNormZero'}):   # Attention for kv_lora
+
+    loras = []
+
+    for _m, _n, _child_module in _find_modules(
+        model,
+        target_replace_module,
+        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
+    ):
+        loras.append((_child_module.lora_up, _child_module.lora_down))
+
+    if len(loras) == 0:
+        raise ValueError("No lora injected.")
+
+    return loras
+
+def save_lora_weight(
+    model,
+    path="./lora.pt",
+    target_replace_module={'AdaLayerNormZero'},     # Attention for kv_lora
+    save_half:bool=False
+):
+    weights = []
+    for _up, _down in extract_lora_ups_down(
+        model, target_replace_module=target_replace_module
+    ):
+        dtype = torch.float16 if save_half else torch.float32
+        weights.append(_up.weight.to("cpu").to(dtype))
+        weights.append(_down.weight.to("cpu").to(dtype))
+
+    torch.save(weights, path)

coz/osediff_sd3.py

@@ -0,0 +1,913 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+import yaml
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Tuple, Optional
+import numpy as np
+import lpips
+from torchvision import transforms
+from PIL import Image
+from peft import LoraConfig, get_peft_model
+
+from copy import deepcopy
+from tqdm import tqdm
+
+from diffusers import StableDiffusion3Pipeline
+from lora.lora_layers import LoraInjectedLinear, LoraInjectedConv2d
+
+from utils.vaehook import VAEHook
+
+
+def inject_lora_vae(vae, lora_rank=4, init_lora_weights="gaussian", verbose=False):
+    """
+    Inject LoRA into the VAE's encoder
+    """
+    vae.requires_grad_(False)
+    vae.train() 
+
+    # Identify modules to LoRA-ify in the encoder
+    l_grep = ["conv1", "conv2", "conv_in", "conv_shortcut", 
+              "conv", "conv_out", "to_k", "to_q", "to_v", "to_out.0"]
+    l_target_modules_encoder = []
+    for n, p in vae.named_parameters():
+        if "bias" in n or "norm" in n:
+            continue
+        for pattern in l_grep:
+            if (pattern in n) and ("encoder" in n):
+                l_target_modules_encoder.append(n.replace(".weight", ""))
+            elif ("quant_conv" in n) and ("post_quant_conv" not in n):
+                l_target_modules_encoder.append(n.replace(".weight", ""))
+
+    if verbose:
+        print("The following VAE parameters will get LoRA:")
+        print(l_target_modules_encoder)
+
+    # Create and add a LoRA adapter
+    lora_conf_encoder = LoraConfig(
+        r=lora_rank,
+        init_lora_weights=init_lora_weights,
+        target_modules=l_target_modules_encoder
+    )
+    
+    adapter_name = "default_encoder"
+    try:
+        vae.add_adapter(lora_conf_encoder, adapter_name=adapter_name)
+        vae.set_adapter(adapter_name)
+    except ValueError as e:
+        if "already exists" in str(e):
+            print(f"Adapter with name {adapter_name} already exists. Skipping injection.")
+        else:
+            raise e
+
+    return vae, l_target_modules_encoder
+
+def _find_modules(model, ancestor_class=None, search_class=[nn.Linear], exclude_children_of=[LoraInjectedLinear]):
+    # Get the targets we should replace all linears under
+    if ancestor_class is not None:
+        ancestors = (
+            module
+            for module in model.modules()
+            if module.__class__.__name__ in ancestor_class
+        )
+    else:
+        # this, in case you want to naively iterate over all modules.
+        ancestors = [module for module in model.modules()]
+
+    for ancestor in ancestors:
+        for fullname, module in ancestor.named_modules():
+            if any([isinstance(module, _class) for _class in search_class]):
+                *path, name = fullname.split(".")
+                parent = ancestor
+                while path:
+                    parent = parent.get_submodule(path.pop(0))
+                if exclude_children_of and any(
+                    [isinstance(parent, _class) for _class in exclude_children_of]
+                ):
+                    continue
+                yield parent, name, module
+
+def inject_lora(model, ancestor_class, loras=None, r:int=4, dropout_p:float=0.0, scale:float=1.0, verbose:bool=False):
+    
+    model.requires_grad_(False)
+    model.train()
+    
+    names = []
+    require_grad_params = []  # to be updated
+
+    total_lora_params = 0  
+
+    if loras is not None:
+        loras = torch.load(loras, map_location=model.device, weights_only=True)
+        loras = [lora.float() for lora in loras]
+
+    for _module, name, _child_module in _find_modules(model, ancestor_class):  # SiLU + Linear Block
+        weight = _child_module.weight
+        bias = _child_module.bias
+
+        if verbose:
+            print(f'LoRA Injection : injecting lora into {name}')
+
+        _tmp = LoraInjectedLinear(
+            _child_module.in_features,
+            _child_module.out_features,
+            _child_module.bias is not None,
+            r=r,
+            dropout_p=dropout_p,
+            scale=scale,
+        )
+        _tmp.linear.weight = nn.Parameter(weight.float())
+        if bias is not None:
+            _tmp.linear.bias = nn.Parameter(bias.float())
+        
+        # switch the module
+        _tmp.to(device=_child_module.weight.device, dtype=torch.float) # keep as float / mixed precision
+        _module._modules[name] = _tmp
+
+        require_grad_params.append(_module._modules[name].lora_up.parameters())
+        require_grad_params.append(_module._modules[name].lora_down.parameters())
+
+        if loras != None:
+            _module._modules[name].lora_up.weight = nn.Parameter(loras.pop(0))
+            _module._modules[name].lora_down.weight = nn.Parameter(loras.pop(0))
+
+        _module._modules[name].lora_up.weight.requires_grad = True
+        _module._modules[name].lora_down.weight.requires_grad = True
+        names.append(name)
+
+        if verbose:
+            # -------- Count LoRA parameters just added --------
+            lora_up_count = sum(p.numel() for p in _tmp.lora_up.parameters())
+            lora_down_count = sum(p.numel() for p in _tmp.lora_down.parameters())
+            lora_total_for_this_layer = lora_up_count + lora_down_count
+            total_lora_params += lora_total_for_this_layer
+            print(f"  Added {lora_total_for_this_layer} params "
+                  f"(lora_up={lora_up_count}, lora_down={lora_down_count})")
+    
+    if verbose:
+        print(f"Total new LoRA parameters added: {total_lora_params}")
+    
+    return require_grad_params, names
+
+def add_mp_hook(transformer):
+    '''
+    For mixed precision of LoRA. (i.e. keep LoRA as float and others as half)
+    '''
+    def pre_hook(module, input):
+        return input.float()
+    
+    def post_hook(module, input, output):
+        return output.half()
+
+    hooks = []
+    for _module, name, _child_module in _find_modules(transformer):
+        if isinstance(_child_module, LoraInjectedLinear):
+            hook = _child_module.lora_up.register_forward_pre_hook(pre_hook)
+            hooks.append(hook)
+            hook = _child_module.lora_down.register_forward_hook(post_hook)
+            hooks.append(hook)
+    
+    return transformer, hooks
+
+def compute_density_for_timestep_sampling(
+    weighting_scheme: str, batch_size: int, logit_mean: float = 0.0, logit_std: float = 1.0, mode_scale: Optional[float] = None
+):
+    """
+    Compute the density for sampling the timesteps when doing SD3 training.
+
+    Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+    SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
+    """
+    if weighting_scheme == "logit_normal":
+        # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
+        u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu")
+        u = torch.nn.functional.sigmoid(u)
+    elif weighting_scheme == "mode":
+        u = torch.rand(size=(batch_size,), device="cpu")
+        u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
+    else:
+        u = torch.rand(size=(batch_size,), device="cpu")
+    return u
+
+def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas):
+    """
+    Computes loss weighting scheme for SD3 training.
+
+    Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+    SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
+    """
+    if weighting_scheme == "sigma_sqrt":
+        weighting = (sigmas**-2.0).float()
+    elif weighting_scheme == "cosmap":
+        bot = 1 - 2 * sigmas + 2 * sigmas**2
+        weighting = 2 / (math.pi * bot)
+    else:
+        weighting = torch.ones_like(sigmas)
+    return weighting
+
+
+class StableDiffusion3Base():
+    def __init__(self, model_key:str='stabilityai/stable-diffusion-3-medium-diffusers', device='cuda', dtype=torch.float16):
+        self.device = device
+        self.dtype = dtype
+
+        pipe = StableDiffusion3Pipeline.from_pretrained(model_key, torch_dtype=self.dtype)
+
+        self.scheduler = pipe.scheduler
+
+        self.tokenizer_1 = pipe.tokenizer
+        self.tokenizer_2 = pipe.tokenizer_2
+        self.tokenizer_3 = pipe.tokenizer_3
+        self.text_enc_1 = pipe.text_encoder.to(device)
+        self.text_enc_2 = pipe.text_encoder_2.to(device)
+        self.text_enc_3 = pipe.text_encoder_3.to(device)
+
+        self.vae=pipe.vae.to(device)
+
+        self.transformer = pipe.transformer.to(device)
+        self.transformer.eval()
+        self.transformer.requires_grad_(False)
+
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels)-1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+
+        del pipe
+
+    def encode_prompt(self, prompt: List[str], batch_size:int=1) -> List[torch.Tensor]:
+        '''
+        We assume that
+        1. number of tokens < max_length
+        2. one prompt for one image
+        '''
+        # CLIP encode (used for modulation of adaLN-zero)
+        # now, we have two CLIPs
+        text_clip1_ids = self.tokenizer_1(prompt,
+                                          padding="max_length",
+                                          max_length=77,
+                                          truncation=True,
+                                          return_tensors='pt').input_ids
+        text_clip1_emb = self.text_enc_1(text_clip1_ids.to(self.device), output_hidden_states=True)
+        pool_clip1_emb = text_clip1_emb[0].to(dtype=self.dtype, device=self.device)
+        text_clip1_emb = text_clip1_emb.hidden_states[-2].to(dtype=self.dtype, device=self.device)
+
+        text_clip2_ids = self.tokenizer_2(prompt,
+                                          padding="max_length",
+                                          max_length=77,
+                                          truncation=True,
+                                          return_tensors='pt').input_ids
+        text_clip2_emb = self.text_enc_2(text_clip2_ids.to(self.device), output_hidden_states=True)
+        pool_clip2_emb = text_clip2_emb[0].to(dtype=self.dtype, device=self.device)
+        text_clip2_emb = text_clip2_emb.hidden_states[-2].to(dtype=self.dtype, device=self.device)
+        
+        # T5 encode (used for text condition)
+        text_t5_ids = self.tokenizer_3(prompt,
+                                       padding="max_length",
+                                       max_length=512,
+                                       truncation=True,
+                                       add_special_tokens=True,
+                                       return_tensors='pt').input_ids
+        text_t5_emb = self.text_enc_3(text_t5_ids.to(self.device))[0]
+        text_t5_emb = text_t5_emb.to(dtype=self.dtype, device=self.device)
+
+        # Merge
+        clip_prompt_emb = torch.cat([text_clip1_emb, text_clip2_emb], dim=-1)
+        clip_prompt_emb = torch.nn.functional.pad(
+            clip_prompt_emb, (0, text_t5_emb.shape[-1] - clip_prompt_emb.shape[-1])
+        )
+        prompt_emb = torch.cat([clip_prompt_emb, text_t5_emb], dim=-2)
+        pooled_prompt_emb = torch.cat([pool_clip1_emb, pool_clip2_emb], dim=-1)
+
+        return prompt_emb, pooled_prompt_emb
+
+    def initialize_latent(self, img_size:Tuple[int], batch_size:int=1, **kwargs):
+        H, W = img_size
+        lH, lW = H//self.vae_scale_factor, W//self.vae_scale_factor
+        lC = self.transformer.config.in_channels
+        latent_shape = (batch_size, lC, lH, lW)
+
+        z = torch.randn(latent_shape, device=self.device, dtype=self.dtype)
+
+        return z
+
+    def encode(self, image: torch.Tensor) -> torch.Tensor:
+        z = self.vae.encode(image).latent_dist.sample()
+        z = (z-self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        return z
+    
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        z = (z/self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        return self.vae.decode(z, return_dict=False)[0]
+    
+    
+class SD3Euler(StableDiffusion3Base):
+    def __init__(self, model_key:str='stabilityai/stable-diffusion-3-medium-diffusers', device='cuda'):
+        super().__init__(model_key=model_key, device=device)
+
+    def inversion(self, src_img, prompts: List[str], NFE:int, cfg_scale: float=1.0, batch_size: int=1):
+
+        # encode text prompts
+        prompt_emb, pooled_emb = self.encode_prompt(prompts, batch_size)
+        null_prompt_emb, null_pooled_emb = self.encode_prompt([""], batch_size)
+
+        # initialize latent
+        src_img = src_img.to(device=self.device, dtype=self.dtype)
+        with torch.no_grad():
+            z = self.encode(src_img)
+            z0 = z.clone()
+
+        # timesteps (default option. You can make your custom here.)
+        self.scheduler.set_timesteps(NFE, device=self.device)
+        timesteps = self.scheduler.timesteps
+        timesteps = torch.cat([timesteps, torch.zeros(1, device=self.device)])
+        timesteps = reversed(timesteps)
+        sigmas = timesteps / self.scheduler.config.num_train_timesteps
+
+        # Solve ODE
+        pbar = tqdm(timesteps[:-1], total=NFE, desc='SD3 Euler Inversion')
+        for i, t in enumerate(pbar):
+            timestep = t.expand(z.shape[0]).to(self.device)
+            pred_v = self.predict_vector(z, timestep, prompt_emb, pooled_emb)
+            if cfg_scale != 1.0:
+                pred_null_v = self.predict_vector(z, timestep, null_prompt_emb, null_pooled_emb)
+            else:
+                pred_null_v = 0.0
+
+            sigma = sigmas[i]
+            sigma_next = sigmas[i+1]
+
+            z = z + (sigma_next - sigma) * (pred_null_v + cfg_scale * (pred_v - pred_null_v))
+
+        return z
+
+    def sample(self, prompts: List[str], NFE:int, img_shape: Optional[Tuple[int]]=None, cfg_scale: float=1.0, batch_size: int = 1, latent:Optional[torch.Tensor]=None):
+        imgH, imgW = img_shape if img_shape is not None else (512, 512)
+
+        # encode text prompts
+        with torch.no_grad():
+            prompt_emb, pooled_emb = self.encode_prompt(prompts, batch_size)
+            null_prompt_emb, null_pooled_emb = self.encode_prompt([""], batch_size)
+
+        # initialize latent
+        if latent is None:
+            z = self.initialize_latent((imgH, imgW), batch_size)
+        else:
+            z = latent
+
+        # timesteps (default option. You can make your custom here.)
+        self.scheduler.set_timesteps(NFE, device=self.device)
+        timesteps = self.scheduler.timesteps
+        sigmas = timesteps / self.scheduler.config.num_train_timesteps
+
+        # Solve ODE
+        pbar = tqdm(timesteps, total=NFE, desc='SD3 Euler')
+        for i, t in enumerate(pbar):
+            timestep = t.expand(z.shape[0]).to(self.device)
+            pred_v = self.predict_vector(z, timestep, prompt_emb, pooled_emb)
+            if cfg_scale != 1.0:
+                pred_null_v = self.predict_vector(z, timestep, null_prompt_emb, null_pooled_emb)
+            else:
+                pred_null_v = 0.0
+
+            sigma = sigmas[i]
+            sigma_next = sigmas[i+1] if i+1 < NFE else 0.0
+
+            z = z + (sigma_next - sigma) * (pred_null_v + cfg_scale * (pred_v - pred_null_v))
+
+        # decode
+        with torch.no_grad():
+            img = self.decode(z)
+        return img
+
+
+class OSEDiff_SD3_GEN(torch.nn.Module):
+    def __init__(self, args, base_model):
+        super().__init__()
+        
+        self.args = args
+        self.model = base_model
+        
+        # Add lora to transformer
+        print('Adding LoRA to OSEDiff_SD3_GEN')
+        self.transformer_gen = copy.deepcopy(self.model.transformer)
+        self.transformer_gen.to('cuda:1')
+        
+        self.transformer_gen.requires_grad_(False)
+        self.transformer_gen.train()
+        self.transformer_gen, hooks = add_mp_hook(self.transformer_gen)
+        self.hooks = hooks
+
+        lora_params, _ = inject_lora(self.transformer_gen, {"AdaLayerNormZero"}, r=args.lora_rank, verbose=False)
+        for name, param in self.transformer_gen.named_parameters():
+            if "lora_" in name:
+                param.requires_grad = True   # LoRA up/down
+            else:
+                param.requires_grad = False  # everything else
+        
+        # Insert LoRA into VAE
+        print("Adding LoRA to VAE")
+        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
+
+    def predict_vector(self, z, t, prompt_emb, pooled_emb):
+        v = self.transformer_gen(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+
+    def forward(self, x_src, batch=None, args=None):
+        
+        z_src = self.model.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device))
+        z_src = z_src.to(self.transformer_gen.device)
+        
+        # calculate prompt_embeddings and neg_prompt_embeddings
+        batch_size, _, _, _ = x_src.shape
+        with torch.no_grad():
+            prompt_embeds, pooled_embeds = self.model.encode_prompt(batch["prompt"], batch_size)
+            neg_prompt_embeds, neg_pooled_embeds = self.model.encode_prompt(batch["neg_prompt"], batch_size)
+        
+        NFE = 1
+        self.model.scheduler.set_timesteps(NFE, device=self.model.device)
+        timesteps = self.model.scheduler.timesteps
+        sigmas = timesteps / self.model.scheduler.config.num_train_timesteps
+        sigmas = sigmas.to(self.transformer_gen.device)
+
+        # Solve ODE
+        i = 0
+        t = timesteps[0]
+
+        timestep = t.expand(z_src.shape[0]).to(self.transformer_gen.device)
+        prompt_embeds = prompt_embeds.to(self.transformer_gen.device, dtype=torch.float32)
+        pooled_embeds = pooled_embeds.to(self.transformer_gen.device, dtype=torch.float32)
+        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
+        pred_null_v = 0.0
+
+        sigma = sigmas[i]
+        sigma_next = sigmas[i+1] if i+1 < NFE else 0.0
+        
+        z_src = z_src + (sigma_next - sigma) * (pred_null_v + 1 * (pred_v - pred_null_v))
+        
+        output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
+
+        return output_image, z_src, prompt_embeds, pooled_embeds
+
+
+class OSEDiff_SD3_REG(torch.nn.Module):
+    def __init__(self, args, base_model):
+        super().__init__()
+        
+        self.args = args
+        self.model = base_model
+        self.transformer_org = self.model.transformer
+        
+        # Add lora to transformer
+        print('Adding LoRA to OSEDiff_SD3_REG')
+        self.transformer_reg = copy.deepcopy(self.transformer_org)
+        self.transformer_reg.to('cuda:1')
+        
+        self.transformer_reg.requires_grad_(False)
+        self.transformer_reg.train()
+        self.transformer_reg, hooks = add_mp_hook(self.transformer_reg)
+        self.hooks = hooks
+        
+        lora_params, _ = inject_lora(self.transformer_reg, {"AdaLayerNormZero"}, r=args.lora_rank, verbose=False)
+        for name, param in self.transformer_reg.named_parameters():
+            if "lora_" in name:
+                param.requires_grad = True   # LoRA up/down
+            else:
+                param.requires_grad = False  # everything else
+        
+    def predict_vector_reg(self, z, t, prompt_emb, pooled_emb):
+        v = self.transformer_reg(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+    
+    def predict_vector_org(self, z, t, prompt_emb, pooled_emb):
+        v = self.transformer_org(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+
+    def distribution_matching_loss(self, z0, prompt_embeds, pooled_embeds, global_step, args):
+        
+        with torch.no_grad():
+            device = self.transformer_reg.device
+            # get timesteps and sigma
+            u = compute_density_for_timestep_sampling(
+                        weighting_scheme="uniform",
+                        batch_size=1,
+                        logit_mean=0.0,
+                        logit_std=1.0,
+                        mode_scale=1.29,
+                    )
+            
+            t_idx = (u*1000).long().to(device)
+            self.model.scheduler.set_timesteps(1000, device=device)
+            times = self.model.scheduler.timesteps
+            t = times[t_idx]
+            sigma = t / 1000
+            
+            # get noise and xt
+            z0 = z0.to(device)
+            noise = torch.randn_like(z0)
+            sigma = sigma.half()
+            zt = (1-sigma) * z0 + sigma * noise
+            
+            # Get x0_prediction of transformer_reg
+            v_pred_reg = self.predict_vector_reg(zt, t, prompt_embeds.to(device), pooled_embeds.to(device))
+            reg_model_pred = v_pred_reg * (-sigma) + zt         # this is x0_prediction for reg
+            
+            # Get x0_prediction of transformer_org
+            org_device = self.transformer_org.device
+            v_pred_org = self.predict_vector_org(zt.to(org_device), t.to(org_device), prompt_embeds.to(org_device), pooled_embeds.to(org_device))
+            org_model_pred = v_pred_org * (-sigma.to(org_device)) + zt.to(org_device)         # this is x0_prediction for org
+            
+            # Visualization
+            if global_step % 100 == 1:
+                self.vsd_visualization(z0, noise, zt, reg_model_pred, org_model_pred, global_step, args)
+            
+        weighting_factor = torch.abs(z0 - org_model_pred.to(device)).mean(dim=[1, 2, 3], keepdim=True)
+
+        grad = (reg_model_pred - org_model_pred.to(device)) / weighting_factor
+        loss = F.mse_loss(z0, (z0 - grad).detach())
+
+        return loss
+    
+    def vsd_visualization(self, z0, noise, zt, reg_model_pred, org_model_pred, global_step, args):
+        #-------- Visualization --------#
+        # 1. Visualize latents, noise, zt
+        z0_img = self.model.decode(z0.to(dtype=torch.float32, device=self.model.vae.device))
+        ns_img = self.model.decode(noise.to(dtype=torch.float32, device=self.model.vae.device))
+        zt_img = self.model.decode(zt.to(dtype=torch.float32, device=self.model.vae.device))
+        
+        z0_img_pil = transforms.ToPILImage()(torch.clamp(z0_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
+        ns_img_pil = transforms.ToPILImage()(torch.clamp(ns_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
+        zt_img_pil = transforms.ToPILImage()(torch.clamp(zt_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
+        
+        # 2. Visualize reg_img, org_img
+        reg_img = self.model.decode(reg_model_pred.to(dtype=torch.float32, device=self.model.vae.device))
+        org_img = self.model.decode(org_model_pred.to(dtype=torch.float32, device=self.model.vae.device))
+        
+        reg_img_pil = transforms.ToPILImage()(torch.clamp(reg_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
+        org_img_pil = transforms.ToPILImage()(torch.clamp(org_img[0].cpu(), -1.0, 1.0) * 0.5 + 0.5)
+        
+        # Concatenate images side by side
+        w, h = z0_img_pil.width, z0_img_pil.height
+        combined_image = Image.new('RGB', (w*5, h))
+        combined_image.paste(z0_img_pil, (0, 0))
+        combined_image.paste(ns_img_pil, (w, 0))
+        combined_image.paste(zt_img_pil, (w*2, 0))
+        combined_image.paste(reg_img_pil, (w*3, 0))
+        combined_image.paste(org_img_pil, (w*4, 0))
+        combined_image.save(os.path.join(args.output_dir, f'visualization/vsd/{global_step}.png'))
+        #-------- Visualization --------#
+    
+    def diff_loss(self, z0, prompt_embeds, pooled_embeds, net_lpips, args):
+        
+        device = self.transformer_reg.device
+        u = compute_density_for_timestep_sampling(
+                        weighting_scheme="uniform",
+                        batch_size=1,
+                        logit_mean=0.0,
+                        logit_std=1.0,
+                        mode_scale=1.29,
+                    )
+        
+        t_idx = (u*1000).long().to(device)
+        self.model.scheduler.set_timesteps(1000, device=device)
+        times = self.model.scheduler.timesteps
+        t = times[t_idx]
+        sigma = t / 1000
+        
+        z0 = z0.to(device)
+        z0, prompt_embeds = z0.detach(), prompt_embeds.detach()
+        noise = torch.randn_like(z0)
+        sigma = sigma.half()
+        zt = (1-sigma) * z0 + sigma * noise         # noisy latents
+
+        # v-prediction
+        v_pred = self.predict_vector_reg(zt, t, prompt_embeds.to(device), pooled_embeds.to(device))
+        model_pred = v_pred * (-sigma) + zt
+        target = z0
+        
+        loss_weight = compute_loss_weighting_for_sd3("logit_normal", sigma)
+        diffusion_loss = loss_weight.float() * F.mse_loss(model_pred.float(), target.float()) 
+        
+        loss_d = diffusion_loss
+
+        return loss_d.mean()
+
+class OSEDiff_SD3_TEST(torch.nn.Module):
+    def __init__(self, args, base_model):
+        super().__init__()
+        
+        self.args = args
+        self.model = base_model
+        self.lora_path = args.lora_path
+        self.vae_path = args.vae_path
+        
+        # Add lora to transformer
+        print(f'Loading LoRA to Transformer from {self.lora_path}')
+        self.model.transformer.requires_grad_(False)
+        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
+        for name, param in self.model.transformer.named_parameters():
+            param.requires_grad = False
+        
+        # Insert LoRA into VAE
+        print(f"Loading LoRA to VAE from {self.vae_path}")
+        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
+        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu")
+        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
+
+    def predict_vector(self, z, t, prompt_emb, pooled_emb):
+        v = self.model.transformer(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+
+    @torch.no_grad()
+    def forward(self, x_src, prompt):
+        
+        z_src = self.model.vae.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device)).latent_dist.sample() * self.model.vae.config.scaling_factor
+        
+        z_src = z_src.to(self.model.transformer.device)
+        
+        # calculate prompt_embeddings and neg_prompt_embeddings
+        batch_size, _, _, _ = x_src.shape
+        with torch.no_grad():
+            prompt_embeds, pooled_embeds = self.model.encode_prompt([prompt], batch_size)
+        
+        self.model.scheduler.set_timesteps(1, device=self.model.device)
+        timesteps = self.model.scheduler.timesteps
+
+        # Solve ODE
+        t = timesteps[0]
+        timestep = t.expand(z_src.shape[0]).to(self.model.transformer.device)
+        prompt_embeds = prompt_embeds.to(self.model.transformer.device, dtype=torch.float32)
+        pooled_embeds = pooled_embeds.to(self.model.transformer.device, dtype=torch.float32)
+        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
+        
+        z_src = z_src - pred_v
+        
+        with torch.no_grad():
+            output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
+
+        return output_image
+
+
+class OSEDiff_SD3_TEST_efficient(torch.nn.Module):
+    def __init__(self, args, base_model):
+        super().__init__()
+        
+        self.args = args
+        self.model = base_model
+        self.lora_path = args.lora_path
+        self.vae_path = args.vae_path
+        
+        # Add lora to transformer
+        print(f'Loading LoRA to Transformer from {self.lora_path}')
+        self.model.transformer.requires_grad_(False)
+        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
+        for name, param in self.model.transformer.named_parameters():
+            param.requires_grad = False
+        
+        # Insert LoRA into VAE
+        print(f"Loading LoRA to VAE from {self.vae_path}")
+        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
+        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu")
+        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
+
+    def predict_vector(self, z, t, prompt_emb, pooled_emb):
+        v = self.model.transformer(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+
+    @torch.no_grad()
+    def forward(self, x_src, prompt):
+        
+        z_src = self.model.vae.encode(x_src.to(dtype=torch.float32, device=self.model.vae.device)).latent_dist.sample() * self.model.vae.config.scaling_factor
+        
+        z_src = z_src.to(self.model.transformer.device)
+        
+        # calculate prompt_embeddings
+        batch_size, _, _, _ = x_src.shape
+        prompt_embeds, pooled_embeds = self.model.encode_prompt([prompt], batch_size)
+        
+        self.model.scheduler.set_timesteps(1, device=self.model.device)
+        timesteps = self.model.scheduler.timesteps
+
+        # Solve ODE
+        t = timesteps[0]
+        timestep = t.expand(z_src.shape[0]).to(self.model.transformer.device)
+        prompt_embeds = prompt_embeds.to(self.model.transformer.device, dtype=torch.float32)
+        pooled_embeds = pooled_embeds.to(self.model.transformer.device, dtype=torch.float32)
+        pred_v = self.predict_vector(z_src, timestep, prompt_embeds, pooled_embeds)
+        z_src = z_src - pred_v
+        
+        output_image = self.model.decode(z_src.to(dtype=torch.float32, device=self.model.vae.device))
+
+        return output_image
+
+
+class OSEDiff_SD3_TEST_TILE(torch.nn.Module):
+    def __init__(self, args, base_model):
+        super().__init__()
+        
+        self.args = args
+        self.model = base_model
+        self.lora_path = args.lora_path
+        self.vae_path = args.vae_path
+        
+        # Add lora to transformer
+        print(f'Loading LoRA to Transformer from {self.lora_path}')
+        self.model.transformer.requires_grad_(False)
+        lora_params, _ = inject_lora(self.model.transformer, {"AdaLayerNormZero"}, loras=self.lora_path, r=args.lora_rank, verbose=False)
+        for name, param in self.model.transformer.named_parameters():
+            param.requires_grad = False
+        
+        # Insert LoRA into VAE
+        print(f"Loading LoRA to VAE from {self.vae_path}")
+        self.model.vae, self.lora_vae_modules_encoder = inject_lora_vae(self.model.vae, lora_rank=args.lora_rank, verbose=False)
+        encoder_state_dict_fp16 = torch.load(self.vae_path, map_location="cpu")
+        self.model.vae.encoder.load_state_dict(encoder_state_dict_fp16)
+        
+        # save original forward (only once)
+        if not hasattr(self.model.vae.encoder, 'original_forward'):
+            setattr(self.model.vae.encoder, 'original_forward', self.model.vae.encoder.forward)
+        if not hasattr(self.model.vae.decoder, 'original_forward'):
+            setattr(self.model.vae.decoder, 'original_forward', self.model.vae.decoder.forward)
+        encoder_tile = args.vae_encoder_tiled_size
+        decoder_tile = args.vae_decoder_tiled_size
+        self.model.vae.encoder.forward = VAEHook(
+            self.model.vae.encoder,
+            tile_size=encoder_tile,
+            is_decoder=False,
+            fast_decoder=False,
+            fast_encoder=True,
+            color_fix=False,
+            to_gpu=True
+        )
+        self.model.vae.decoder.forward = VAEHook(
+            self.model.vae.decoder,
+            tile_size=decoder_tile,
+            is_decoder=True,
+            fast_decoder=True,
+            fast_encoder=False,
+            color_fix=False,
+            to_gpu=True
+        )
+
+    def predict_vector(self, z, t, prompt_emb, pooled_emb):
+        v = self.model.transformer(hidden_states=z,
+                             timestep=t,
+                             pooled_projections=pooled_emb,
+                             encoder_hidden_states=prompt_emb,
+                             return_dict=False)[0]
+        return v
+    
+    @torch.no_grad()
+    def create_full_latent(self, x_full: torch.Tensor, vlm_model, vlm_processor, full_path, next_path, prompt_type) -> torch.Tensor:
+        device = self.model.transformer.device
+        # 1) encode to full latent (via VAEHook)
+        z_full = self.model.vae.encode(x_full).latent_dist.sample() \
+                * self.model.vae.config.scaling_factor
+        z_full = z_full.to(device)
+        B, C, H, W = z_full.shape
+
+        # 2) grid size
+        tsize = self.args.latent_tiled_size
+        tover  = self.args.latent_tiled_overlap
+        stride = tsize - tover
+        rows = (H - tsize + stride - 1)//stride + 1
+        cols = (W - tsize + stride - 1)//stride + 1
+        print(f'TILE SIZE: {tsize}, TILE OVERLAP: {tover}, STRIDE: {stride}, ROWS: {rows}, COLS: {cols}')
+
+        # 3) make gaussian weight patched [B,C,tsize,tsize]
+        weights = self._make_gaussian(tsize, tsize, 1).to(z_full.device)
+
+        # 4) collect all patches
+        positions = []
+        out_tiles = []
+        timestep = self.model.scheduler.timesteps[0].expand(B).to(z_full.device)
+
+        for i in range(rows):
+            y0 = min(i*stride, H - tsize)
+            for j in range(cols):
+                x0 = min(j*stride, W - tsize)
+                positions.append((y0, x0))
+                patch = z_full[:, :, y0:y0+tsize, x0:x0+tsize]
+                
+                # decode and save patch for later usage
+                patch_path = f'{next_path[:-4]}_patch_row{i}col{j}.png'
+                patch_img = self.decode_full_latent(patch)
+                patch_pil = transforms.ToPILImage()((patch_img[0] * 0.5 + 0.5).clamp(0,1))
+                patch_pil.save(patch_path)
+                
+                # create prompt to explain patch (CoZ)
+                prompt = self.create_prompt(vlm_model, vlm_processor, full_path, patch_path, prompt_type)
+                print('PROMPT: ', prompt)
+                prompt_emb, pooled_emb = self.model.encode_prompt([prompt], batch_size=B)
+                prompt_emb = prompt_emb.to(z_full.device, dtype=torch.float32)
+                pooled_emb = pooled_emb.to(z_full.device, dtype=torch.float32)
+                
+                v = self.predict_vector(patch, timestep, prompt_emb, pooled_emb)
+                out_tiles.append(patch - v)
+
+        # 5) accumulate + normalize
+        z_out = torch.zeros_like(z_full)
+        z_norm = torch.zeros_like(z_full)
+        norm  = torch.zeros_like(z_full)
+
+        for (y0,x0), tile in zip(positions, out_tiles):
+            z_out[:, :, y0:y0+tsize, x0:x0+tsize] += tile
+            z_norm[:, :, y0:y0+tsize, x0:x0+tsize] += tile * weights
+            norm[:, :, y0:y0+tsize, x0:x0+tsize] += weights
+
+        # 6) avoid division by zero and finalize
+        eps = 1e-10
+        z_norm = z_norm / (norm + eps)
+        return z_norm, z_out
+
+    @torch.no_grad()
+    def decode_full_latent(self, z_full: torch.Tensor) -> torch.Tensor:
+        """
+        Decode the tiled full latent into an RGB image (with tiled VAE decoder).
+        """
+        z_full = z_full.to(self.model.vae.device)
+        img = self.model.vae.decode(z_full / self.model.vae.config.scaling_factor).sample
+        return img.clamp(-1,1)
+    
+    @torch.no_grad()
+    def create_prompt(self, vlm_model, vlm_processor, full_path, patch_path, prompt_type):
+        if prompt_type in ('vlm','vlm_base'):
+            from qwen_vl_utils import process_vision_info
+            
+            message_text = None
+            start_image_path = full_path
+            input_image_path = patch_path
+            
+            message_text = "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 words."
+            messages = [
+                {"role": "system", "content": f"{message_text}"},
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": start_image_path},
+                        {"type": "image", "image": input_image_path}
+                    ]
+                }
+            ]
+
+            text = vlm_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            image_inputs, video_inputs = process_vision_info(messages)
+            inputs = vlm_processor(
+                text=[text],
+                images=image_inputs,
+                videos=video_inputs,
+                padding=True,
+                return_tensors="pt",
+            )
+            generated_ids = vlm_model.generate(**inputs, max_new_tokens=16)
+            generated_ids_trimmed = [
+                out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            ]
+            output_text = vlm_processor.batch_decode(
+                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+            )
+
+            prompt_text = output_text[0]
+            return prompt_text
+        else:
+            raise ValueError(f"Unknown prompt_type: {prompt_type}")
+    
+    def _make_gaussian(self, w, h, nb):
+        from numpy import pi, exp, sqrt
+        import numpy as np
+
+        latent_width = w
+        latent_height = h
+
+        var = 0.01
+        midpoint = (latent_width - 1) / 2       # -1 because index goes from 0 to latent_width - 1
+        x_probs = [exp(-(x-midpoint)*(x-midpoint)/(latent_width*latent_width)/(2*var)) / sqrt(2*pi*var) for x in range(latent_width)]
+        midpoint = latent_height / 2
+        y_probs = [exp(-(y-midpoint)*(y-midpoint)/(latent_height*latent_height)/(2*var)) / sqrt(2*pi*var) for y in range(latent_height)]
+
+        weights = np.outer(y_probs, x_probs)
+        return torch.tile(torch.tensor(weights, device=self.model.vae.device), (nb, self.model.vae.config.latent_channels, 1, 1))

coz/utils/devices.py

@@ -0,0 +1,138 @@
+import sys
+import contextlib
+from functools import lru_cache
+
+import torch
+#from modules import errors
+
+if sys.platform == "darwin":
+    from modules import mac_specific
+
+
+def has_mps() -> bool:
+    if sys.platform != "darwin":
+        return False
+    else:
+        return mac_specific.has_mps
+
+
+def get_cuda_device_string():
+    return "cuda"
+
+
+def get_optimal_device_name():
+    if torch.cuda.is_available():
+        return get_cuda_device_string()
+
+    if has_mps():
+        return "mps"
+
+    return "cpu"
+
+
+def get_optimal_device():
+    return torch.device(get_optimal_device_name())
+
+
+def get_device_for(task):
+    return get_optimal_device()
+
+
+def torch_gc():
+
+    if torch.cuda.is_available():
+        with torch.cuda.device(get_cuda_device_string()):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    if has_mps():
+        mac_specific.torch_mps_gc()
+
+
+def enable_tf32():
+    if torch.cuda.is_available():
+
+        # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
+        # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
+        if any(torch.cuda.get_device_capability(devid) == (7, 5) for devid in range(0, torch.cuda.device_count())):
+            torch.backends.cudnn.benchmark = True
+
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+
+enable_tf32()
+#errors.run(enable_tf32, "Enabling TF32")
+
+cpu = torch.device("cpu")
+device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = torch.device("cuda")
+dtype = torch.float16
+dtype_vae = torch.float16
+dtype_unet = torch.float16
+unet_needs_upcast = False
+
+
+def cond_cast_unet(input):
+    return input.to(dtype_unet) if unet_needs_upcast else input
+
+
+def cond_cast_float(input):
+    return input.float() if unet_needs_upcast else input
+
+
+def randn(seed, shape):
+    torch.manual_seed(seed)
+    return torch.randn(shape, device=device)
+
+
+def randn_without_seed(shape):
+    return torch.randn(shape, device=device)
+
+
+def autocast(disable=False):
+    if disable:
+        return contextlib.nullcontext()
+
+    return torch.autocast("cuda")
+
+
+def without_autocast(disable=False):
+    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+
+
+class NansException(Exception):
+    pass
+
+
+def test_for_nans(x, where):
+    if not torch.all(torch.isnan(x)).item():
+        return
+
+    if where == "unet":
+        message = "A tensor with all NaNs was produced in Unet."
+
+    elif where == "vae":
+        message = "A tensor with all NaNs was produced in VAE."
+
+    else:
+        message = "A tensor with all NaNs was produced."
+
+    message += " Use --disable-nan-check commandline argument to disable this check."
+
+    raise NansException(message)
+
+
+@lru_cache
+def first_time_calculation():
+    """
+    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
+    spends about 2.7 seconds doing that, at least wih NVidia.
+    """
+
+    x = torch.zeros((1, 1)).to(device, dtype)
+    linear = torch.nn.Linear(1, 1).to(device, dtype)
+    linear(x)
+
+    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
+    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
+    conv2d(x)

coz/utils/vaehook.py

@@ -0,0 +1,829 @@
+# ------------------------------------------------------------------------
+#
+#   Ultimate VAE Tile Optimization
+#
+#   Introducing a revolutionary new optimization designed to make
+#   the VAE work with giant images on limited VRAM!
+#   Say goodbye to the frustration of OOM and hello to seamless output!
+#
+# ------------------------------------------------------------------------
+#
+#   This script is a wild hack that splits the image into tiles,
+#   encodes each tile separately, and merges the result back together.
+#
+#   Advantages:
+#   - The VAE can now work with giant images on limited VRAM
+#       (~10 GB for 8K images!)
+#   - The merged output is completely seamless without any post-processing.
+#
+#   Drawbacks:
+#   - Giant RAM needed. To store the intermediate results for a 4096x4096
+#       images, you need 32 GB RAM it consumes ~20GB); for 8192x8192
+#       you need 128 GB RAM machine (it consumes ~100 GB)
+#   - NaNs always appear in for 8k images when you use fp16 (half) VAE
+#       You must use --no-half-vae to disable half VAE for that giant image.
+#   - Slow speed. With default tile size, it takes around 50/200 seconds
+#       to encode/decode a 4096x4096 image; and 200/900 seconds to encode/decode
+#       a 8192x8192 image. (The speed is limited by both the GPU and the CPU.)
+#   - The gradient calculation is not compatible with this hack. It
+#       will break any backward() or torch.autograd.grad() that passes VAE.
+#       (But you can still use the VAE to generate training data.)
+#
+#   How it works:
+#   1) The image is split into tiles.
+#       - To ensure perfect results, each tile is padded with 32 pixels
+#           on each side.
+#       - Then the conv2d/silu/upsample/downsample can produce identical
+#           results to the original image without splitting.
+#   2) The original forward is decomposed into a task queue and a task worker.
+#       - The task queue is a list of functions that will be executed in order.
+#       - The task worker is a loop that executes the tasks in the queue.
+#   3) The task queue is executed for each tile.
+#       - Current tile is sent to GPU.
+#       - local operations are directly executed.
+#       - Group norm calculation is temporarily suspended until the mean
+#           and var of all tiles are calculated.
+#       - The residual is pre-calculated and stored and addded back later.
+#       - When need to go to the next tile, the current tile is send to cpu.
+#   4) After all tiles are processed, tiles are merged on cpu and return.
+#
+#   Enjoy!
+#
+#   @author: LI YI @ Nanyang Technological University - Singapore
+#   @date: 2023-03-02
+#   @license: MIT License
+#
+#   Please give me a star if you like this project!
+#
+# -------------------------------------------------------------------------
+
+import gc
+from time import time
+import math
+from tqdm import tqdm
+
+import torch
+import torch.version
+import torch.nn.functional as F
+from einops import rearrange
+import os
+import sys
+sys.path.append(os.getcwd())
+import utils.devices as devices
+
+try:
+    import xformers
+    import xformers.ops
+except ImportError:
+    pass
+
+sd_flag = False
+
+def get_recommend_encoder_tile_size():
+    if torch.cuda.is_available():
+        total_memory = torch.cuda.get_device_properties(
+            devices.device).total_memory // 2**20
+        if total_memory > 16*1000:
+            ENCODER_TILE_SIZE = 3072
+        elif total_memory > 12*1000:
+            ENCODER_TILE_SIZE = 2048
+        elif total_memory > 8*1000:
+            ENCODER_TILE_SIZE = 1536
+        else:
+            ENCODER_TILE_SIZE = 960
+    else:
+        ENCODER_TILE_SIZE = 512
+    return ENCODER_TILE_SIZE
+
+
+def get_recommend_decoder_tile_size():
+    if torch.cuda.is_available():
+        total_memory = torch.cuda.get_device_properties(
+            devices.device).total_memory // 2**20
+        if total_memory > 30*1000:
+            DECODER_TILE_SIZE = 256
+        elif total_memory > 16*1000:
+            DECODER_TILE_SIZE = 192
+        elif total_memory > 12*1000:
+            DECODER_TILE_SIZE = 128
+        elif total_memory > 8*1000:
+            DECODER_TILE_SIZE = 96
+        else:
+            DECODER_TILE_SIZE = 64
+    else:
+        DECODER_TILE_SIZE = 64
+    return DECODER_TILE_SIZE
+
+
+if 'global const':
+    DEFAULT_ENABLED = False
+    DEFAULT_MOVE_TO_GPU = False
+    DEFAULT_FAST_ENCODER = True
+    DEFAULT_FAST_DECODER = True
+    DEFAULT_COLOR_FIX = 0
+    DEFAULT_ENCODER_TILE_SIZE = get_recommend_encoder_tile_size()
+    DEFAULT_DECODER_TILE_SIZE = get_recommend_decoder_tile_size()
+
+
+# inplace version of silu
+def inplace_nonlinearity(x):
+    # Test: fix for Nans
+    return F.silu(x, inplace=True)
+
+# extracted from ldm.modules.diffusionmodules.model
+
+# from diffusers lib
+def attn_forward_new(self, h_):
+    batch_size, channel, height, width = h_.shape
+    hidden_states = h_.view(batch_size, channel, height * width).transpose(1, 2)
+
+    attention_mask = None
+    encoder_hidden_states = None
+    batch_size, sequence_length, _ = hidden_states.shape
+    attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+    query = self.to_q(hidden_states)
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+    elif self.norm_cross:
+        encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
+
+    key = self.to_k(encoder_hidden_states)
+    value = self.to_v(encoder_hidden_states)
+
+    query = self.head_to_batch_dim(query)
+    key = self.head_to_batch_dim(key)
+    value = self.head_to_batch_dim(value)
+
+    attention_probs = self.get_attention_scores(query, key, attention_mask)
+    hidden_states = torch.bmm(attention_probs, value)
+    hidden_states = self.batch_to_head_dim(hidden_states)
+
+    # linear proj
+    hidden_states = self.to_out[0](hidden_states)
+    # dropout
+    hidden_states = self.to_out[1](hidden_states)
+
+    hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+    return hidden_states
+
+def attn_forward(self, h_):
+    q = self.q(h_)
+    k = self.k(h_)
+    v = self.v(h_)
+
+    # compute attention
+    b, c, h, w = q.shape
+    q = q.reshape(b, c, h*w)
+    q = q.permute(0, 2, 1)   # b,hw,c
+    k = k.reshape(b, c, h*w)  # b,c,hw
+    w_ = torch.bmm(q, k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+    w_ = w_ * (int(c)**(-0.5))
+    w_ = torch.nn.functional.softmax(w_, dim=2)
+
+    # attend to values
+    v = v.reshape(b, c, h*w)
+    w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+    # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+    h_ = torch.bmm(v, w_)
+    h_ = h_.reshape(b, c, h, w)
+
+    h_ = self.proj_out(h_)
+
+    return h_
+
+
+def xformer_attn_forward(self, h_):
+    q = self.q(h_)
+    k = self.k(h_)
+    v = self.v(h_)
+
+    # compute attention
+    B, C, H, W = q.shape
+    q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
+
+    q, k, v = map(
+        lambda t: t.unsqueeze(3)
+        .reshape(B, t.shape[1], 1, C)
+        .permute(0, 2, 1, 3)
+        .reshape(B * 1, t.shape[1], C)
+        .contiguous(),
+        (q, k, v),
+    )
+    out = xformers.ops.memory_efficient_attention(
+        q, k, v, attn_bias=None, op=self.attention_op)
+
+    out = (
+        out.unsqueeze(0)
+        .reshape(B, 1, out.shape[1], C)
+        .permute(0, 2, 1, 3)
+        .reshape(B, out.shape[1], C)
+    )
+    out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
+    out = self.proj_out(out)
+    return out
+
+
+def attn2task(task_queue, net):
+    if False: #isinstance(net, AttnBlock):
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.norm))
+        task_queue.append(('attn', lambda x, net=net: attn_forward(net, x)))
+        task_queue.append(['add_res', None])
+    elif False: #isinstance(net, MemoryEfficientAttnBlock):
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.norm))
+        task_queue.append(
+            ('attn', lambda x, net=net: xformer_attn_forward(net, x)))
+        task_queue.append(['add_res', None])
+    else:
+        task_queue.append(('store_res', lambda x: x))
+        task_queue.append(('pre_norm', net.group_norm))
+        task_queue.append(('attn', lambda x, net=net: attn_forward_new(net, x)))
+        task_queue.append(['add_res', None])
+
+def resblock2task(queue, block):
+    """
+    Turn a ResNetBlock into a sequence of tasks and append to the task queue
+
+    @param queue: the target task queue
+    @param block: ResNetBlock
+
+    """
+    if block.in_channels != block.out_channels:
+        if sd_flag:
+            if block.use_conv_shortcut:
+                queue.append(('store_res', block.conv_shortcut))
+            else:
+                queue.append(('store_res', block.nin_shortcut))
+        else:
+            if block.use_in_shortcut:
+                queue.append(('store_res', block.conv_shortcut))
+            else:
+                queue.append(('store_res', block.nin_shortcut))
+
+    else:
+        queue.append(('store_res', lambda x: x))
+    queue.append(('pre_norm', block.norm1))
+    queue.append(('silu', inplace_nonlinearity))
+    queue.append(('conv1', block.conv1))
+    queue.append(('pre_norm', block.norm2))
+    queue.append(('silu', inplace_nonlinearity))
+    queue.append(('conv2', block.conv2))
+    queue.append(['add_res', None])
+
+
+
+def build_sampling(task_queue, net, is_decoder):
+    """
+    Build the sampling part of a task queue
+    @param task_queue: the target task queue
+    @param net: the network
+    @param is_decoder: currently building decoder or encoder
+    """
+    if is_decoder:
+        # resblock2task(task_queue, net.mid.block_1)
+        # attn2task(task_queue, net.mid.attn_1)
+        # resblock2task(task_queue, net.mid.block_2)
+        # resolution_iter = reversed(range(net.num_resolutions))
+        # block_ids = net.num_res_blocks + 1
+        # condition = 0
+        # module = net.up
+        # func_name = 'upsample'
+        resblock2task(task_queue, net.mid_block.resnets[0])
+        attn2task(task_queue, net.mid_block.attentions[0])
+        resblock2task(task_queue, net.mid_block.resnets[1])
+        resolution_iter = (range(len(net.up_blocks)))  # range(0,4)
+        block_ids = 2 + 1
+        condition = len(net.up_blocks) - 1
+        module = net.up_blocks
+        func_name = 'upsamplers'
+    else:
+        # resolution_iter = range(net.num_resolutions)
+        # block_ids = net.num_res_blocks
+        # condition = net.num_resolutions - 1
+        # module = net.down
+        # func_name = 'downsample'
+        resolution_iter = (range(len(net.down_blocks)))  # range(0,4)
+        block_ids = 2
+        condition = len(net.down_blocks) - 1
+        module = net.down_blocks
+        func_name = 'downsamplers'
+
+
+    for i_level in resolution_iter:
+        for i_block in range(block_ids):
+            resblock2task(task_queue, module[i_level].resnets[i_block])
+        if i_level != condition:
+            if is_decoder:
+                task_queue.append((func_name, module[i_level].upsamplers[0]))
+            else:
+                task_queue.append((func_name, module[i_level].downsamplers[0]))
+
+    if not is_decoder:
+        resblock2task(task_queue, net.mid_block.resnets[0])
+        attn2task(task_queue, net.mid_block.attentions[0])
+        resblock2task(task_queue, net.mid_block.resnets[1])
+
+
+def build_task_queue(net, is_decoder):
+    """
+    Build a single task queue for the encoder or decoder
+    @param net: the VAE decoder or encoder network
+    @param is_decoder: currently building decoder or encoder
+    @return: the task queue
+    """
+    task_queue = []
+    task_queue.append(('conv_in', net.conv_in))
+
+    # construct the sampling part of the task queue
+    # because encoder and decoder share the same architecture, we extract the sampling part
+    build_sampling(task_queue, net, is_decoder)
+    if is_decoder and not sd_flag:
+        net.give_pre_end = False
+        net.tanh_out = False
+
+    if not is_decoder or not net.give_pre_end:
+        if sd_flag:
+            task_queue.append(('pre_norm', net.norm_out))
+        else:
+            task_queue.append(('pre_norm', net.conv_norm_out))
+        task_queue.append(('silu', inplace_nonlinearity))
+        task_queue.append(('conv_out', net.conv_out))
+        if is_decoder and net.tanh_out:
+            task_queue.append(('tanh', torch.tanh))
+
+    return task_queue
+
+
+def clone_task_queue(task_queue):
+    """
+    Clone a task queue
+    @param task_queue: the task queue to be cloned
+    @return: the cloned task queue
+    """
+    return [[item for item in task] for task in task_queue]
+
+
+def get_var_mean(input, num_groups, eps=1e-6):
+    """
+    Get mean and var for group norm
+    """
+    b, c = input.size(0), input.size(1)
+    channel_in_group = int(c/num_groups)
+    input_reshaped = input.contiguous().view(
+        1, int(b * num_groups), channel_in_group, *input.size()[2:])
+    var, mean = torch.var_mean(
+        input_reshaped, dim=[0, 2, 3, 4], unbiased=False)
+    return var, mean
+
+
+def custom_group_norm(input, num_groups, mean, var, weight=None, bias=None, eps=1e-6):
+    """
+    Custom group norm with fixed mean and var
+
+    @param input: input tensor
+    @param num_groups: number of groups. by default, num_groups = 32
+    @param mean: mean, must be pre-calculated by get_var_mean
+    @param var: var, must be pre-calculated by get_var_mean
+    @param weight: weight, should be fetched from the original group norm
+    @param bias: bias, should be fetched from the original group norm
+    @param eps: epsilon, by default, eps = 1e-6 to match the original group norm
+
+    @return: normalized tensor
+    """
+    b, c = input.size(0), input.size(1)
+    channel_in_group = int(c/num_groups)
+    input_reshaped = input.contiguous().view(
+        1, int(b * num_groups), channel_in_group, *input.size()[2:])
+
+    out = F.batch_norm(input_reshaped, mean, var, weight=None, bias=None,
+                       training=False, momentum=0, eps=eps)
+
+    out = out.view(b, c, *input.size()[2:])
+
+    # post affine transform
+    if weight is not None:
+        out *= weight.view(1, -1, 1, 1)
+    if bias is not None:
+        out += bias.view(1, -1, 1, 1)
+    return out
+
+
+def crop_valid_region(x, input_bbox, target_bbox, is_decoder):
+    """
+    Crop the valid region from the tile
+    @param x: input tile
+    @param input_bbox: original input bounding box
+    @param target_bbox: output bounding box
+    @param scale: scale factor
+    @return: cropped tile
+    """
+    padded_bbox = [i * 8 if is_decoder else i//8 for i in input_bbox]
+    margin = [target_bbox[i] - padded_bbox[i] for i in range(4)]
+    return x[:, :, margin[2]:x.size(2)+margin[3], margin[0]:x.size(3)+margin[1]]
+
+# ↓↓↓ https://github.com/Kahsolt/stable-diffusion-webui-vae-tile-infer ↓↓↓
+
+
+def perfcount(fn):
+    def wrapper(*args, **kwargs):
+        ts = time()
+
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats(devices.device)
+        devices.torch_gc()
+        gc.collect()
+
+        ret = fn(*args, **kwargs)
+
+        devices.torch_gc()
+        gc.collect()
+        if torch.cuda.is_available():
+            vram = torch.cuda.max_memory_allocated(devices.device) / 2**20
+            torch.cuda.reset_peak_memory_stats(devices.device)
+            print(
+                f'[Tiled VAE]: Done in {time() - ts:.3f}s, max VRAM alloc {vram:.3f} MB')
+        else:
+            print(f'[Tiled VAE]: Done in {time() - ts:.3f}s')
+
+        return ret
+    return wrapper
+
+# copy end :)
+
+
+class GroupNormParam:
+    def __init__(self):
+        self.var_list = []
+        self.mean_list = []
+        self.pixel_list = []
+        self.weight = None
+        self.bias = None
+
+    def add_tile(self, tile, layer):
+        var, mean = get_var_mean(tile, 32)
+        # For giant images, the variance can be larger than max float16
+        # In this case we create a copy to float32
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+        # ============= DEBUG: test for infinite =============
+        # if torch.isinf(var).any():
+        #    print('var: ', var)
+        # ====================================================
+        self.var_list.append(var)
+        self.mean_list.append(mean)
+        self.pixel_list.append(
+            tile.shape[2]*tile.shape[3])
+        if hasattr(layer, 'weight'):
+            self.weight = layer.weight
+            self.bias = layer.bias
+        else:
+            self.weight = None
+            self.bias = None
+
+    def summary(self):
+        """
+        summarize the mean and var and return a function
+        that apply group norm on each tile
+        """
+        if len(self.var_list) == 0:
+            return None
+        var = torch.vstack(self.var_list)
+        mean = torch.vstack(self.mean_list)
+        max_value = max(self.pixel_list)
+        pixels = torch.tensor(
+            self.pixel_list, dtype=torch.float32, device=devices.device) / max_value
+        sum_pixels = torch.sum(pixels)
+        pixels = pixels.unsqueeze(
+            1) / sum_pixels
+        var = torch.sum(
+            var * pixels, dim=0)
+        mean = torch.sum(
+            mean * pixels, dim=0)
+        return lambda x:  custom_group_norm(x, 32, mean, var, self.weight, self.bias)
+
+    @staticmethod
+    def from_tile(tile, norm):
+        """
+        create a function from a single tile without summary
+        """
+        var, mean = get_var_mean(tile, 32)
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+            # if it is a macbook, we need to convert back to float16
+            if var.device.type == 'mps':
+                # clamp to avoid overflow
+                var = torch.clamp(var, 0, 60000)
+                var = var.half()
+                mean = mean.half()
+        if hasattr(norm, 'weight'):
+            weight = norm.weight
+            bias = norm.bias
+        else:
+            weight = None
+            bias = None
+
+        def group_norm_func(x, mean=mean, var=var, weight=weight, bias=bias):
+            return custom_group_norm(x, 32, mean, var, weight, bias, 1e-6)
+        return group_norm_func
+
+
+class VAEHook:
+    def __init__(self, net, tile_size, is_decoder, fast_decoder, fast_encoder, color_fix, to_gpu=False):
+        self.net = net                  # encoder | decoder
+        self.tile_size = tile_size
+        self.is_decoder = is_decoder
+        self.fast_mode = (fast_encoder and not is_decoder) or (
+            fast_decoder and is_decoder)
+        self.color_fix = color_fix and not is_decoder
+        self.to_gpu = to_gpu
+        self.pad = 11 if is_decoder else 32
+
+    def __call__(self, x):
+        B, C, H, W = x.shape
+        original_device = next(self.net.parameters()).device
+        try:
+            if self.to_gpu:
+                # self.net.to(devices.get_optimal_device())
+                self.net.to(original_device)
+            if max(H, W) <= self.pad * 2 + self.tile_size:
+                # print("[Tiled VAE]: the input size is tiny and unnecessary to tile.")
+                return self.net.original_forward(x).to(original_device)
+            else:
+                return self.vae_tile_forward(x)
+        finally:
+            self.net.to(original_device)
+
+    def get_best_tile_size(self, lowerbound, upperbound):
+        """
+        Get the best tile size for GPU memory
+        """
+        divider = 32
+        while divider >= 2:
+            remainer = lowerbound % divider
+            if remainer == 0:
+                return lowerbound
+            candidate = lowerbound - remainer + divider
+            if candidate <= upperbound:
+                return candidate
+            divider //= 2
+        return lowerbound
+
+    def split_tiles(self, h, w):
+        """
+        Tool function to split the image into tiles
+        @param h: height of the image
+        @param w: width of the image
+        @return: tile_input_bboxes, tile_output_bboxes
+        """
+        tile_input_bboxes, tile_output_bboxes = [], []
+        tile_size = self.tile_size
+        pad = self.pad
+        num_height_tiles = math.ceil((h - 2 * pad) / tile_size)
+        num_width_tiles = math.ceil((w - 2 * pad) / tile_size)
+        # If any of the numbers are 0, we let it be 1
+        # This is to deal with long and thin images
+        num_height_tiles = max(num_height_tiles, 1)
+        num_width_tiles = max(num_width_tiles, 1)
+
+        # Suggestions from https://github.com/Kahsolt: auto shrink the tile size
+        real_tile_height = math.ceil((h - 2 * pad) / num_height_tiles)
+        real_tile_width = math.ceil((w - 2 * pad) / num_width_tiles)
+        real_tile_height = self.get_best_tile_size(real_tile_height, tile_size)
+        real_tile_width = self.get_best_tile_size(real_tile_width, tile_size)
+
+        print(f'[Tiled VAE]: split to {num_height_tiles}x{num_width_tiles} = {num_height_tiles*num_width_tiles} tiles. ' +
+              f'Optimal tile size {real_tile_width}x{real_tile_height}, original tile size {tile_size}x{tile_size}')
+
+        for i in range(num_height_tiles):
+            for j in range(num_width_tiles):
+                # bbox: [x1, x2, y1, y2]
+                # the padding is is unnessary for image borders. So we directly start from (32, 32)
+                input_bbox = [
+                    pad + j * real_tile_width,
+                    min(pad + (j + 1) * real_tile_width, w),
+                    pad + i * real_tile_height,
+                    min(pad + (i + 1) * real_tile_height, h),
+                ]
+
+                # if the output bbox is close to the image boundary, we extend it to the image boundary
+                output_bbox = [
+                    input_bbox[0] if input_bbox[0] > pad else 0,
+                    input_bbox[1] if input_bbox[1] < w - pad else w,
+                    input_bbox[2] if input_bbox[2] > pad else 0,
+                    input_bbox[3] if input_bbox[3] < h - pad else h,
+                ]
+
+                # scale to get the final output bbox
+                output_bbox = [x * 8 if self.is_decoder else x // 8 for x in output_bbox]
+                tile_output_bboxes.append(output_bbox)
+
+                # indistinguishable expand the input bbox by pad pixels
+                tile_input_bboxes.append([
+                    max(0, input_bbox[0] - pad),
+                    min(w, input_bbox[1] + pad),
+                    max(0, input_bbox[2] - pad),
+                    min(h, input_bbox[3] + pad),
+                ])
+
+        return tile_input_bboxes, tile_output_bboxes
+
+    @torch.no_grad()
+    def estimate_group_norm(self, z, task_queue, color_fix):
+        device = z.device
+        tile = z
+        last_id = len(task_queue) - 1
+        while last_id >= 0 and task_queue[last_id][0] != 'pre_norm':
+            last_id -= 1
+        if last_id <= 0 or task_queue[last_id][0] != 'pre_norm':
+            raise ValueError('No group norm found in the task queue')
+        # estimate until the last group norm
+        for i in range(last_id + 1):
+            task = task_queue[i]
+            if task[0] == 'pre_norm':
+                group_norm_func = GroupNormParam.from_tile(tile, task[1])
+                task_queue[i] = ('apply_norm', group_norm_func)
+                if i == last_id:
+                    return True
+                tile = group_norm_func(tile)
+            elif task[0] == 'store_res':
+                task_id = i + 1
+                while task_id < last_id and task_queue[task_id][0] != 'add_res':
+                    task_id += 1
+                if task_id >= last_id:
+                    continue
+                task_queue[task_id][1] = task[1](tile)
+            elif task[0] == 'add_res':
+                tile += task[1].to(device)
+                task[1] = None
+            elif color_fix and task[0] == 'downsample':
+                for j in range(i, last_id + 1):
+                    if task_queue[j][0] == 'store_res':
+                        task_queue[j] = ('store_res_cpu', task_queue[j][1])
+                return True
+            else:
+                tile = task[1](tile)
+            try:
+                devices.test_for_nans(tile, "vae")
+            except:
+                print(f'Nan detected in fast mode estimation. Fast mode disabled.')
+                return False
+
+        raise IndexError('Should not reach here')
+
+    # @perfcount
+    @torch.no_grad()
+    def vae_tile_forward(self, z):
+        """
+        Decode a latent vector z into an image in a tiled manner.
+        @param z: latent vector
+        @return: image
+        """
+        device = next(self.net.parameters()).device
+        net = self.net
+        tile_size = self.tile_size
+        is_decoder = self.is_decoder
+
+        z = z.detach() # detach the input to avoid backprop
+
+        N, height, width = z.shape[0], z.shape[2], z.shape[3]
+        net.last_z_shape = z.shape
+
+        # Split the input into tiles and build a task queue for each tile
+        print(f'[Tiled VAE]: input_size: {z.shape}, tile_size: {tile_size}, padding: {self.pad}')
+
+        in_bboxes, out_bboxes = self.split_tiles(height, width)
+
+        # Prepare tiles by split the input latents
+        tiles = []
+        for input_bbox in in_bboxes:
+            tile = z[:, :, input_bbox[2]:input_bbox[3], input_bbox[0]:input_bbox[1]].cpu()
+            tiles.append(tile)
+
+        num_tiles = len(tiles)
+        num_completed = 0
+
+        # Build task queues
+        single_task_queue = build_task_queue(net, is_decoder)
+        #print(single_task_queue)
+        if self.fast_mode:
+            # Fast mode: downsample the input image to the tile size,
+            # then estimate the group norm parameters on the downsampled image
+            scale_factor = tile_size / max(height, width)
+            z = z.to(device)
+            downsampled_z = F.interpolate(z, scale_factor=scale_factor, mode='nearest-exact')
+            # use nearest-exact to keep statictics as close as possible
+            print(f'[Tiled VAE]: Fast mode enabled, estimating group norm parameters on {downsampled_z.shape[3]} x {downsampled_z.shape[2]} image')
+
+            # ======= Special thanks to @Kahsolt for distribution shift issue ======= #
+            # The downsampling will heavily distort its mean and std, so we need to recover it.
+            std_old, mean_old = torch.std_mean(z, dim=[0, 2, 3], keepdim=True)
+            std_new, mean_new = torch.std_mean(downsampled_z, dim=[0, 2, 3], keepdim=True)
+            downsampled_z = (downsampled_z - mean_new) / std_new * std_old + mean_old
+            del std_old, mean_old, std_new, mean_new
+            # occasionally the std_new is too small or too large, which exceeds the range of float16
+            # so we need to clamp it to max z's range.
+            downsampled_z = torch.clamp_(downsampled_z, min=z.min(), max=z.max())
+            estimate_task_queue = clone_task_queue(single_task_queue)
+            if self.estimate_group_norm(downsampled_z, estimate_task_queue, color_fix=self.color_fix):
+                single_task_queue = estimate_task_queue
+            del downsampled_z
+
+        task_queues = [clone_task_queue(single_task_queue) for _ in range(num_tiles)]
+
+        # Dummy result
+        result = None
+        result_approx = None
+        #try:
+        #    with devices.autocast():
+        #        result_approx = torch.cat([F.interpolate(cheap_approximation(x).unsqueeze(0), scale_factor=opt_f, mode='nearest-exact') for x in z], dim=0).cpu()
+        #except: pass
+        # Free memory of input latent tensor
+        del z
+
+        # Task queue execution
+        pbar = tqdm(total=num_tiles * len(task_queues[0]), desc=f"[Tiled VAE]: Executing {'Decoder' if is_decoder else 'Encoder'} Task Queue: ")
+
+        # execute the task back and forth when switch tiles so that we always
+        # keep one tile on the GPU to reduce unnecessary data transfer
+        forward = True
+        interrupted = False
+        #state.interrupted = interrupted
+        while True:
+            #if state.interrupted: interrupted = True ; break
+
+            group_norm_param = GroupNormParam()
+            for i in range(num_tiles) if forward else reversed(range(num_tiles)):
+                #if state.interrupted: interrupted = True ; break
+
+                tile = tiles[i].to(device)
+                input_bbox = in_bboxes[i]
+                task_queue = task_queues[i]
+
+                interrupted = False
+                while len(task_queue) > 0:
+                    #if state.interrupted: interrupted = True ; break
+
+                    # DEBUG: current task
+                    # print('Running task: ', task_queue[0][0], ' on tile ', i, '/', num_tiles, ' with shape ', tile.shape)
+                    task = task_queue.pop(0)
+                    if task[0] == 'pre_norm':
+                        group_norm_param.add_tile(tile, task[1])
+                        break
+                    elif task[0] == 'store_res' or task[0] == 'store_res_cpu':
+                        task_id = 0
+                        res = task[1](tile)
+                        if not self.fast_mode or task[0] == 'store_res_cpu':
+                            res = res.cpu()
+                        while task_queue[task_id][0] != 'add_res':
+                            task_id += 1
+                        task_queue[task_id][1] = res
+                    elif task[0] == 'add_res':
+                        tile += task[1].to(device)
+                        task[1] = None
+                    else:
+                        tile = task[1](tile)
+                    pbar.update(1)
+
+                if interrupted: break
+
+                # check for NaNs in the tile.
+                # If there are NaNs, we abort the process to save user's time
+                #devices.test_for_nans(tile, "vae")
+
+                #print(tiles[i].shape, tile.shape, i, num_tiles)
+                if len(task_queue) == 0:
+                    tiles[i] = None
+                    num_completed += 1
+                    if result is None:      # NOTE: dim C varies from different cases, can only be inited dynamically
+                        result = torch.zeros((N, tile.shape[1], height * 8 if is_decoder else height // 8, width * 8 if is_decoder else width // 8), device=device, requires_grad=False)
+                    result[:, :, out_bboxes[i][2]:out_bboxes[i][3], out_bboxes[i][0]:out_bboxes[i][1]] = crop_valid_region(tile, in_bboxes[i], out_bboxes[i], is_decoder)
+                    del tile
+                elif i == num_tiles - 1 and forward:
+                    forward = False
+                    tiles[i] = tile
+                elif i == 0 and not forward:
+                    forward = True
+                    tiles[i] = tile
+                else:
+                    tiles[i] = tile.cpu()
+                    del tile
+
+            if interrupted: break
+            if num_completed == num_tiles: break
+
+            # insert the group norm task to the head of each task queue
+            group_norm_func = group_norm_param.summary()
+            if group_norm_func is not None:
+                for i in range(num_tiles):
+                    task_queue = task_queues[i]
+                    task_queue.insert(0, ('apply_norm', group_norm_func))
+
+        # Done!
+        pbar.close()
+        return result if result is not None else result_approx.to(device)

coz/utils/wavelet_color_fix.py

@@ -0,0 +1,119 @@
+'''
+# --------------------------------------------------------------------------------
+#   Color fixed script from Li Yi (https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py)
+# --------------------------------------------------------------------------------
+'''
+
+import torch
+from PIL import Image
+from torch import Tensor
+from torch.nn import functional as F
+
+from torchvision.transforms import ToTensor, ToPILImage
+
+def adain_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply adaptive instance normalization
+    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def wavelet_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply wavelet reconstruction
+    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def calc_mean_std(feat: Tensor, eps=1e-5):
+    """Calculate mean and std for adaptive_instance_normalization.
+    Args:
+        feat (Tensor): 4D tensor.
+        eps (float): A small value added to the variance to avoid
+            divide-by-zero. Default: 1e-5.
+    """
+    size = feat.size()
+    assert len(size) == 4, 'The input feature should be 4D tensor.'
+    b, c = size[:2]
+    feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
+    feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
+    return feat_mean, feat_std
+
+def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
+    """Adaptive instance normalization.
+    Adjust the reference features to have the similar color and illuminations
+    as those in the degradate features.
+    Args:
+        content_feat (Tensor): The reference feature.
+        style_feat (Tensor): The degradate features.
+    """
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+def wavelet_blur(image: Tensor, radius: int):
+    """
+    Apply wavelet blur to the input tensor.
+    """
+    # input shape: (1, 3, H, W)
+    # convolution kernel
+    kernel_vals = [
+        [0.0625, 0.125, 0.0625],
+        [0.125, 0.25, 0.125],
+        [0.0625, 0.125, 0.0625],
+    ]
+    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
+    # add channel dimensions to the kernel to make it a 4D tensor
+    kernel = kernel[None, None]
+    # repeat the kernel across all input channels
+    kernel = kernel.repeat(3, 1, 1, 1)
+    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
+    # apply convolution
+    output = F.conv2d(image, kernel, groups=3, dilation=radius)
+    return output
+
+def wavelet_decomposition(image: Tensor, levels=5):
+    """
+    Apply wavelet decomposition to the input tensor.
+    This function only returns the low frequency & the high frequency.
+    """
+    high_freq = torch.zeros_like(image)
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur(image, radius)
+        high_freq += (image - low_freq)
+        image = low_freq
+
+    return high_freq, low_freq
+
+def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
+    """
+    Apply wavelet decomposition, so that the content will have the same color as the style.
+    """
+    # calculate the wavelet decomposition of the content feature
+    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
+    del content_low_freq
+    # calculate the wavelet decomposition of the style feature
+    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
+    del style_high_freq
+    # reconstruct the content feature with the style's high frequency
+    return content_high_freq + style_low_freq

inference.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""OracleZoom - faithful extreme super-resolution (4x -> 16x -> 64x -> 256x).
+
+Self-contained: this repo + auto-downloaded Stable Diffusion 3-medium and Qwen2.5-VL-3B.
+Vendored Chain-of-Zoom code lives in ./coz, checkpoints in ./ckpt, merged model = merged_transformer.safetensors.
+
+Usage:
+    python inference.py --input ./inputs --output ./outputs
+Outputs: outputs/per-scale/scale1..4/<name>.png  (4x / 16x / 64x / 256x)
+"""
+import argparse, glob, os, sys
+import torch
+from PIL import Image
+from torchvision import transforms
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(HERE, "coz"))   # vendored Chain-of-Zoom modules
+
+COZ_PROMPT = ("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 words.")
+_to_tensor = transforms.Compose([transforms.ToTensor()])
+
+
+def resize_and_center_crop(img, size):
+    w, h = img.size
+    scale = size / min(w, h)
+    nw, nh = int(w * scale), int(h * scale)
+    img = img.resize((nw, nh), Image.LANCZOS)
+    l, t = (nw - size) // 2, (nh - size) // 2
+    return img.crop((l, t, l + size, t + size))
+
+
+class _SRArgs:
+    def __init__(self, ckpt, sd3, process_size):
+        self.lora_path = f"{ckpt}/SR_LoRA/model_20001.pkl"
+        self.vae_path = f"{ckpt}/SR_VAE/vae_encoder_20001.pt"
+        self.pretrained_model_name_or_path = sd3
+        self.process_size = process_size
+        self.lora_rank = 4
+        self.merge_and_unload_lora = False
+        self.mixed_precision = "fp16"
+
+
+def build_sr(ckpt, sd3, process_size):
+    from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler
+    sr = SD3Euler()
+    for m in [sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae]:
+        m.to("cuda:0")
+    sr.transformer.to("cuda:0", dtype=torch.float32)
+    sr.vae.to("cuda:0", dtype=torch.float32)
+    for m in [sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae]:
+        m.requires_grad_(False)
+    return OSEDiff_SD3_TEST(_SRArgs(ckpt, sd3, process_size), sr)
+
+
+def build_vlm(ckpt):
+    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+    from qwen_vl_utils import process_vision_info
+    from peft import PeftModel
+    name = "Qwen/Qwen2.5-VL-3B-Instruct"
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        name, torch_dtype="auto", device_map="auto", attn_implementation="sdpa")
+    proc = AutoProcessor.from_pretrained(name)
+    model = PeftModel.from_pretrained(model, f"{ckpt}/VLM_LoRA/checkpoint-10000").merge_and_unload().eval()
+    return model, proc, process_vision_info
+
+
+def vlm_prompt(model, proc, pvi, first, second, max_new_tokens=32):
+    messages = [{"role": "system", "content": COZ_PROMPT},
+                {"role": "user", "content": [{"type": "image", "image": first},
+                                             {"type": "image", "image": second}]}]
+    text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    ii, vi = pvi(messages)
+    inputs = proc(text=[text], images=ii, videos=vi, padding=True, return_tensors="pt").to("cuda")
+    gen = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    trimmed = [o[len(i):] for i, o in zip(inputs.input_ids, gen)]
+    return proc.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--input", required=True, help="folder of input images")
+    ap.add_argument("--output", required=True, help="output folder")
+    ap.add_argument("--merged", default=os.path.join(HERE, "merged_transformer.safetensors"))
+    ap.add_argument("--ckpt", default=os.path.join(HERE, "ckpt"))
+    ap.add_argument("--sd3", default="stabilityai/stable-diffusion-3-medium-diffusers")
+    ap.add_argument("--rec_num", type=int, default=4)
+    ap.add_argument("--upscale", type=int, default=4)
+    ap.add_argument("--process_size", type=int, default=512)
+    ap.add_argument("--max_new_tokens", type=int, default=32)
+    a = ap.parse_args()
+    os.makedirs(a.output, exist_ok=True)
+
+    sr = build_sr(a.ckpt, a.sd3, a.process_size)
+    from safetensors.torch import load_file
+    sd = load_file(a.merged)
+    dev = next(sr.model.transformer.parameters()).device
+    sd = {k: v.to(dev, dtype=torch.float32) for k, v in sd.items()}
+    miss, unexp = sr.model.transformer.load_state_dict(sd, strict=False)
+    print(f"[OracleZoom] merged transformer loaded (missing={len(miss)} unexpected={len(unexp)})", flush=True)
+    model, proc, pvi = build_vlm(a.ckpt)
+
+    imgs = sorted(p for e in ("*.png", "*.jpg", "*.jpeg", "*.webp") for p in glob.glob(f"{a.input}/{e}"))
+    for img_path in imgs:
+        bname = os.path.splitext(os.path.basename(img_path))[0]
+        rec_dir = os.path.join(a.output, "per-sample", bname)
+        os.makedirs(rec_dir, exist_ok=True)
+        cur = resize_and_center_crop(Image.open(img_path).convert("RGB"), a.process_size)
+        cur.save(f"{rec_dir}/0.png")
+        for rec in range(a.rec_num):
+            prev = Image.open(f"{rec_dir}/{rec}.png").convert("RGB")
+            w, h = prev.size
+            nw, nh = w // a.upscale, h // a.upscale
+            crop = prev.crop(((w - nw) // 2, (h - nh) // 2, (w + nw) // 2, (h + nh) // 2))
+            zoom = crop.resize((w, h), Image.BICUBIC)
+            zp = f"{rec_dir}/{rec + 1}_input.png"
+            zoom.save(zp)
+            prompt = vlm_prompt(model, proc, pvi, f"{rec_dir}/{rec}.png", zp, a.max_new_tokens)
+            lq = _to_tensor(zoom).unsqueeze(0).to("cuda") * 2 - 1
+            with torch.no_grad():
+                out = torch.clamp(sr(lq, prompt=prompt)[0].cpu(), -1.0, 1.0)
+            transforms.ToPILImage()(out * 0.5 + 0.5).save(f"{rec_dir}/{rec + 1}.png")
+            print(f"  {bname}  scale{rec + 1} ({4 ** (rec + 1)}x): {prompt}", flush=True)
+        for s in range(a.rec_num + 1):
+            d = os.path.join(a.output, "per-scale", f"scale{s}")
+            os.makedirs(d, exist_ok=True)
+            Image.open(f"{rec_dir}/{s}.png").save(os.path.join(d, f"{bname}.png"))
+    print("[OracleZoom] done ->", a.output, flush=True)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,16 @@
+torch==2.4.1
+torchvision
+diffusers==0.32.1
+transformers==4.49.0
+peft
+accelerate
+safetensors
+huggingface_hub
+qwen-vl-utils
+pyyaml
+lpips
+einops
+numpy<2
+pillow
+sentencepiece
+protobuf