Spaces:
Paused
Paused
| import importlib | |
| from functools import partial | |
| import cv2 | |
| import numpy as np | |
| import safetensors | |
| import safetensors.torch | |
| import torch | |
| import torch.nn as nn | |
| from inspect import isfunction | |
| from omegaconf import OmegaConf | |
| from src.smplfusion import DDIM, share, scheduler | |
| from .common import * | |
| DOWNLOAD_URL = 'https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/resolve/main/x4-upscaler-ema.safetensors?download=true' | |
| MODEL_PATH = f'{MODEL_FOLDER}/sd-2-0-upsample/x4-upscaler-ema.safetensors' | |
| # pre-download | |
| # download_file(DOWNLOAD_URL, MODEL_PATH) | |
| def exists(x): | |
| return x is not None | |
| def default(val, d): | |
| if exists(val): | |
| return val | |
| return d() if isfunction(d) else d | |
| def extract_into_tensor(a, t, x_shape): | |
| b, *_ = t.shape | |
| out = a.gather(-1, t) | |
| return out.reshape(b, *((1,) * (len(x_shape) - 1))) | |
| def predict_eps_from_z_and_v(schedule, x_t, t, v): | |
| return ( | |
| extract_into_tensor(schedule.sqrt_alphas.to(x_t.device), t, x_t.shape) * v + | |
| extract_into_tensor(schedule.sqrt_one_minus_alphas.to(x_t.device), t, x_t.shape) * x_t | |
| ) | |
| def predict_start_from_z_and_v(schedule, x_t, t, v): | |
| return ( | |
| extract_into_tensor(schedule.sqrt_alphas.to(x_t.device), t, x_t.shape) * x_t - | |
| extract_into_tensor(schedule.sqrt_one_minus_alphas.to(x_t.device), t, x_t.shape) * v | |
| ) | |
| def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): | |
| if schedule == "linear": | |
| betas = ( | |
| torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 | |
| ) | |
| elif schedule == "cosine": | |
| timesteps = ( | |
| torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s | |
| ) | |
| alphas = timesteps / (1 + cosine_s) * np.pi / 2 | |
| alphas = torch.cos(alphas).pow(2) | |
| alphas = alphas / alphas[0] | |
| betas = 1 - alphas[1:] / alphas[:-1] | |
| betas = np.clip(betas, a_min=0, a_max=0.999) | |
| elif schedule == "sqrt_linear": | |
| betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) | |
| elif schedule == "sqrt": | |
| betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 | |
| else: | |
| raise ValueError(f"schedule '{schedule}' unknown.") | |
| return betas.numpy() | |
| def disabled_train(self, mode=True): | |
| """Overwrite model.train with this function to make sure train/eval mode | |
| does not change anymore.""" | |
| return self | |
| class AbstractLowScaleModel(nn.Module): | |
| # for concatenating a downsampled image to the latent representation | |
| def __init__(self, noise_schedule_config=None): | |
| super(AbstractLowScaleModel, self).__init__() | |
| if noise_schedule_config is not None: | |
| self.register_schedule(**noise_schedule_config) | |
| def register_schedule(self, beta_schedule="linear", timesteps=1000, | |
| linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): | |
| betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, | |
| cosine_s=cosine_s) | |
| alphas = 1. - betas | |
| alphas_cumprod = np.cumprod(alphas, axis=0) | |
| alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) | |
| timesteps, = betas.shape | |
| self.num_timesteps = int(timesteps) | |
| self.linear_start = linear_start | |
| self.linear_end = linear_end | |
| assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' | |
| to_torch = partial(torch.tensor, dtype=torch.float32) | |
| self.register_buffer('betas', to_torch(betas)) | |
| self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) | |
| self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) | |
| # calculations for diffusion q(x_t | x_{t-1}) and others | |
| self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) | |
| self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) | |
| self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) | |
| self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) | |
| self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) | |
| def q_sample(self, x_start, t, noise=None): | |
| noise = default(noise, lambda: torch.randn_like(x_start)) | |
| return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + | |
| extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) | |
| def forward(self, x): | |
| return x, None | |
| def decode(self, x): | |
| return x | |
| class ImageConcatWithNoiseAugmentation(AbstractLowScaleModel): | |
| def __init__(self, noise_schedule_config, max_noise_level=1000, to_cuda=False): | |
| super().__init__(noise_schedule_config=noise_schedule_config) | |
| self.max_noise_level = max_noise_level | |
| def forward(self, x, noise_level=None): | |
| if noise_level is None: | |
| noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long() | |
| else: | |
| assert isinstance(noise_level, torch.Tensor) | |
| z = self.q_sample(x, noise_level) | |
| return z, noise_level | |
| def get_obj_from_str(string): | |
| module, cls = string.rsplit(".", 1) | |
| try: | |
| return getattr(importlib.import_module(module, package=None), cls) | |
| except: | |
| return getattr(importlib.import_module('src.' + module, package=None), cls) | |
| def load_obj(path): | |
| objyaml = OmegaConf.load(path) | |
| return get_obj_from_str(objyaml['__class__'])(**objyaml.get("__init__", {})) | |
| def load_model(dtype=torch.bfloat16, device='cuda:0'): | |
| download_file(DOWNLOAD_URL, MODEL_PATH) | |
| state_dict = safetensors.torch.load_file(MODEL_PATH) | |
| config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v2-upsample.yaml') | |
| unet = load_obj(f'{CONFIG_FOLDER}/unet/upsample/v2.yaml').eval().cuda() | |
| vae = load_obj(f'{CONFIG_FOLDER}/vae-upsample.yaml').eval().cuda() | |
| encoder = load_obj(f'{CONFIG_FOLDER}/encoders/openclip.yaml').eval().cuda() | |
| ddim = DDIM(config, vae, encoder, unet) | |
| extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x} | |
| unet_state = extract(state_dict, 'model.diffusion_model') | |
| encoder_state = extract(state_dict, 'cond_stage_model') | |
| vae_state = extract(state_dict, 'first_stage_model') | |
| unet.load_state_dict(unet_state) | |
| encoder.load_state_dict(encoder_state) | |
| vae.load_state_dict(vae_state) | |
| unet = unet.requires_grad_(False) | |
| encoder = encoder.requires_grad_(False) | |
| vae = vae.requires_grad_(False) | |
| unet.to(dtype=dtype, device=device) | |
| vae.to(dtype=dtype, device=device) | |
| encoder.to(dtype=dtype, device=device) | |
| encoder.device = device | |
| ddim = DDIM(config, vae, encoder, unet) | |
| params = { | |
| 'noise_schedule_config': { | |
| 'linear_start': 0.0001, | |
| 'linear_end': 0.02 | |
| }, | |
| 'max_noise_level': 350 | |
| } | |
| low_scale_model = ImageConcatWithNoiseAugmentation(**params).eval().to('cuda') | |
| low_scale_model.train = disabled_train | |
| for param in low_scale_model.parameters(): | |
| param.requires_grad = False | |
| low_scale_model = low_scale_model.to(dtype=dtype, device=device) | |
| ddim.low_scale_model = low_scale_model | |
| return ddim | |