benos/CogVideoX-5B-Space

1import math2from typing import Union, List3 4import torch5import os6from datetime import datetime7import numpy as np8import itertools9import PIL.Image10import safetensors.torch11import tqdm12import logging13from diffusers.utils import export_to_video14from spandrel import ModelLoader15 16logger = logging.getLogger(__file__)17 18 19def load_torch_file(ckpt, device=None, dtype=torch.float16):20    if device is None:21        device = torch.device("cpu")22    if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):23        sd = safetensors.torch.load_file(ckpt, device=device.type)24    else:25        if not "weights_only" in torch.load.__code__.co_varnames:26            logger.warning(27                "Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely."28            )29 30        pl_sd = torch.load(ckpt, map_location=device, weights_only=True)31        if "global_step" in pl_sd:32            logger.debug(f"Global Step: {pl_sd['global_step']}")33        if "state_dict" in pl_sd:34            sd = pl_sd["state_dict"]35        elif "params_ema" in pl_sd:36            sd = pl_sd["params_ema"]37        else:38            sd = pl_sd39 40    sd = {k: v.to(dtype) for k, v in sd.items()}41    return sd42 43 44def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False):45    if filter_keys:46        out = {}47    else:48        out = state_dict49    for rp in replace_prefix:50        replace = list(51            map(52                lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp) :])),53                filter(lambda a: a.startswith(rp), state_dict.keys()),54            )55        )56        for x in replace:57            w = state_dict.pop(x[0])58            out[x[1]] = w59    return out60 61 62def module_size(module):63    module_mem = 064    sd = module.state_dict()65    for k in sd:66        t = sd[k]67        module_mem += t.nelement() * t.element_size()68    return module_mem69 70 71def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):72    return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap)))73 74 75@torch.inference_mode()76def tiled_scale_multidim(77    samples, function, tile=(64, 64), overlap=8, upscale_amount=4, out_channels=3, output_device="cpu", pbar=None78):79    dims = len(tile)80    print(f"samples dtype:{samples.dtype}")81    output = torch.empty(82        [samples.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), samples.shape[2:])),83        device=output_device,84    )85 86    for b in range(samples.shape[0]):87        s = samples[b : b + 1]88        out = torch.zeros(89            [s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])),90            device=output_device,91        )92        out_div = torch.zeros(93            [s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])),94            device=output_device,95        )96 97        for it in itertools.product(*map(lambda a: range(0, a[0], a[1] - overlap), zip(s.shape[2:], tile))):98            s_in = s99            upscaled = []100 101            for d in range(dims):102                pos = max(0, min(s.shape[d + 2] - overlap, it[d]))103                l = min(tile[d], s.shape[d + 2] - pos)104                s_in = s_in.narrow(d + 2, pos, l)105                upscaled.append(round(pos * upscale_amount))106 107            ps = function(s_in).to(output_device)108            mask = torch.ones_like(ps)109            feather = round(overlap * upscale_amount)110            for t in range(feather):111                for d in range(2, dims + 2):112                    m = mask.narrow(d, t, 1)113                    m *= (1.0 / feather) * (t + 1)114                    m = mask.narrow(d, mask.shape[d] - 1 - t, 1)115                    m *= (1.0 / feather) * (t + 1)116 117            o = out118            o_d = out_div119            for d in range(dims):120                o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2])121                o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2])122 123            o += ps * mask124            o_d += mask125 126            if pbar is not None:127                pbar.update(1)128 129        output[b : b + 1] = out / out_div130    return output131 132 133def tiled_scale(134    samples,135    function,136    tile_x=64,137    tile_y=64,138    overlap=8,139    upscale_amount=4,140    out_channels=3,141    output_device="cpu",142    pbar=None,143):144    return tiled_scale_multidim(145        samples, function, (tile_y, tile_x), overlap, upscale_amount, out_channels, output_device, pbar146    )147 148 149def load_sd_upscale(ckpt, inf_device):150    sd = load_torch_file(ckpt, device=inf_device)151    if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:152        sd = state_dict_prefix_replace(sd, {"module.": ""})153    out = ModelLoader().load_from_state_dict(sd).half()154    return out155 156 157def upscale(upscale_model, tensor: torch.Tensor, inf_device, output_device="cpu") -> torch.Tensor:158    memory_required = module_size(upscale_model.model)159    memory_required += (160        (512 * 512 * 3) * tensor.element_size() * max(upscale_model.scale, 1.0) * 384.0161    )  # The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate162    memory_required += tensor.nelement() * tensor.element_size()163    print(f"UPScaleMemory required: {memory_required / 1024 / 1024 / 1024} GB")164 165    upscale_model.to(inf_device)166    tile = 512167    overlap = 32168 169    steps = tensor.shape[0] * get_tiled_scale_steps(170        tensor.shape[3], tensor.shape[2], tile_x=tile, tile_y=tile, overlap=overlap171    )172 173    pbar = ProgressBar(steps, desc="Tiling and Upscaling")174 175    s = tiled_scale(176        samples=tensor.to(torch.float16),177        function=lambda a: upscale_model(a),178        tile_x=tile,179        tile_y=tile,180        overlap=overlap,181        upscale_amount=upscale_model.scale,182        pbar=pbar,183    )184 185    upscale_model.to(output_device)186    return s187 188 189def upscale_batch_and_concatenate(upscale_model, latents, inf_device, output_device="cpu") -> torch.Tensor:190    upscaled_latents = []191    for i in range(latents.size(0)):192        latent = latents[i]193        upscaled_latent = upscale(upscale_model, latent, inf_device, output_device)194        upscaled_latents.append(upscaled_latent)195    return torch.stack(upscaled_latents)196 197 198def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8):199    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")200    video_path = f"./output/{timestamp}.mp4"201    os.makedirs(os.path.dirname(video_path), exist_ok=True)202    export_to_video(tensor, video_path, fps=fps)203    return video_path204 205 206class ProgressBar:207    def __init__(self, total, desc=None):208        self.total = total209        self.current = 0210        self.b_unit = tqdm.tqdm(total=total, desc="ProgressBar context index: 0" if desc is None else desc)211 212    def update(self, value):213        if value > self.total:214            value = self.total215        self.current = value216        if self.b_unit is not None:217            self.b_unit.set_description("ProgressBar context index: {}".format(self.current))218            self.b_unit.refresh()219 220            # 更新进度221            self.b_unit.update(self.current)222