Add files via upload
ADD
This commit is contained in:
smthemex
@@ -0,0 +1,45 @@
|
||||
from typing import Optional
|
||||
import random
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class SimilarImageFilter:
|
||||
def __init__(self, threshold: float = 0.98, max_skip_frame: float = 10) -> None:
|
||||
self.threshold = threshold
|
||||
self.prev_tensor = None
|
||||
self.cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6)
|
||||
self.max_skip_frame = max_skip_frame
|
||||
self.skip_count = 0
|
||||
|
||||
def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]:
|
||||
if self.prev_tensor is None:
|
||||
self.prev_tensor = x.detach().clone()
|
||||
return x
|
||||
else:
|
||||
cos_sim = self.cos(self.prev_tensor.reshape(-1), x.reshape(-1)).item()
|
||||
sample = random.uniform(0, 1)
|
||||
if self.threshold >= 1:
|
||||
skip_prob = 0
|
||||
else:
|
||||
skip_prob = max(0, 1 - (1 - cos_sim) / (1 - self.threshold))
|
||||
|
||||
# not skip frame
|
||||
if skip_prob < sample:
|
||||
self.prev_tensor = x.detach().clone()
|
||||
return x
|
||||
# skip frame
|
||||
else:
|
||||
if self.skip_count > self.max_skip_frame:
|
||||
self.skip_count = 0
|
||||
self.prev_tensor = x.detach().clone()
|
||||
return x
|
||||
else:
|
||||
self.skip_count += 1
|
||||
return None
|
||||
|
||||
def set_threshold(self, threshold: float) -> None:
|
||||
self.threshold = threshold
|
||||
|
||||
def set_max_skip_frame(self, max_skip_frame: float) -> None:
|
||||
self.max_skip_frame = max_skip_frame
|
||||
@@ -0,0 +1,173 @@
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torchvision
|
||||
|
||||
|
||||
def denormalize(images: Union[torch.Tensor, np.ndarray]) -> torch.Tensor:
|
||||
"""
|
||||
Denormalize an image array to [0,1].
|
||||
"""
|
||||
return (images / 2 + 0.5).clamp(0, 1)
|
||||
|
||||
|
||||
def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
|
||||
"""
|
||||
Convert a PyTorch tensor to a NumPy image.
|
||||
"""
|
||||
images = images.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return images
|
||||
|
||||
|
||||
def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
|
||||
"""
|
||||
Convert a NumPy image or a batch of images to a PIL image.
|
||||
"""
|
||||
if images.ndim == 3:
|
||||
images = images[None, ...]
|
||||
images = (images * 255).round().astype("uint8")
|
||||
if images.shape[-1] == 1:
|
||||
# special case for grayscale (single channel) images
|
||||
pil_images = [
|
||||
PIL.Image.fromarray(image.squeeze(), mode="L") for image in images
|
||||
]
|
||||
else:
|
||||
pil_images = [PIL.Image.fromarray(image) for image in images]
|
||||
|
||||
return pil_images
|
||||
|
||||
|
||||
def postprocess_image(
|
||||
image: torch.Tensor,
|
||||
output_type: str = "pil",
|
||||
do_denormalize: Optional[List[bool]] = None,
|
||||
) -> Union[torch.Tensor, np.ndarray, PIL.Image.Image]:
|
||||
if not isinstance(image, torch.Tensor):
|
||||
raise ValueError(
|
||||
f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
|
||||
)
|
||||
|
||||
if output_type == "latent":
|
||||
return image
|
||||
|
||||
do_normalize_flg = True
|
||||
if do_denormalize is None:
|
||||
do_denormalize = [do_normalize_flg] * image.shape[0]
|
||||
|
||||
image = torch.stack(
|
||||
[
|
||||
denormalize(image[i]) if do_denormalize[i] else image[i]
|
||||
for i in range(image.shape[0])
|
||||
]
|
||||
)
|
||||
|
||||
if output_type == "pt":
|
||||
return image
|
||||
|
||||
image = pt_to_numpy(image)
|
||||
|
||||
if output_type == "np":
|
||||
return image
|
||||
|
||||
if output_type == "pil":
|
||||
return numpy_to_pil(image)
|
||||
|
||||
|
||||
def process_image(
|
||||
image_pil: PIL.Image.Image, range: Tuple[int, int] = (-1, 1)
|
||||
) -> Tuple[torch.Tensor, PIL.Image.Image]:
|
||||
image = torchvision.transforms.ToTensor()(image_pil)
|
||||
r_min, r_max = range[0], range[1]
|
||||
image = image * (r_max - r_min) + r_min
|
||||
return image[None, ...], image_pil
|
||||
|
||||
|
||||
def pil2tensor(image_pil: PIL.Image.Image) -> torch.Tensor:
|
||||
height = image_pil.height
|
||||
width = image_pil.width
|
||||
imgs = []
|
||||
img, _ = process_image(image_pil)
|
||||
imgs.append(img)
|
||||
imgs = torch.vstack(imgs)
|
||||
images = torch.nn.functional.interpolate(
|
||||
imgs, size=(height, width), mode="bilinear"
|
||||
)
|
||||
image_tensors = images.to(torch.float16)
|
||||
return image_tensors
|
||||
|
||||
### Optical flow utils
|
||||
|
||||
def coords_grid(b, h, w, homogeneous=False, device=None):
|
||||
y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) # [H, W]
|
||||
|
||||
stacks = [x, y]
|
||||
|
||||
if homogeneous:
|
||||
ones = torch.ones_like(x) # [H, W]
|
||||
stacks.append(ones)
|
||||
|
||||
grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W]
|
||||
|
||||
grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W]
|
||||
|
||||
if device is not None:
|
||||
grid = grid.to(device)
|
||||
|
||||
return grid
|
||||
|
||||
def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
|
||||
b, c, h, w = feature.size()
|
||||
assert flow.size(1) == 2
|
||||
|
||||
grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W]
|
||||
|
||||
return bilinear_sample(feature, grid, padding_mode=padding_mode,
|
||||
return_mask=mask)
|
||||
|
||||
def bilinear_sample(img, sample_coords, mode='bilinear', padding_mode='zeros', return_mask=False):
|
||||
# img: [B, C, H, W]
|
||||
# sample_coords: [B, 2, H, W] in image scale
|
||||
if sample_coords.size(1) != 2: # [B, H, W, 2]
|
||||
sample_coords = sample_coords.permute(0, 3, 1, 2)
|
||||
|
||||
b, _, h, w = sample_coords.shape
|
||||
|
||||
# Normalize to [-1, 1]
|
||||
x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
|
||||
y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
|
||||
|
||||
grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2]
|
||||
|
||||
img = torch.nn.functional.grid_sample(img, grid, mode=mode, padding_mode=padding_mode, align_corners=True)
|
||||
|
||||
if return_mask:
|
||||
mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) # [B, H, W]
|
||||
|
||||
return img, mask
|
||||
|
||||
return img
|
||||
|
||||
def forward_backward_consistency_check(fwd_flow, bwd_flow,
|
||||
alpha=0.1,
|
||||
beta=0.5
|
||||
):
|
||||
# fwd_flow, bwd_flow: [B, 2, H, W]
|
||||
# alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
|
||||
assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
|
||||
assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
|
||||
flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W]
|
||||
|
||||
warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
|
||||
warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
|
||||
|
||||
diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
|
||||
diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
|
||||
|
||||
threshold = alpha * flow_mag + beta
|
||||
|
||||
fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
|
||||
bwd_occ = (diff_bwd > threshold).float()
|
||||
|
||||
return fwd_occ, bwd_occ
|
||||
@@ -0,0 +1,52 @@
|
||||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Dict, Optional
|
||||
|
||||
from packaging.version import Version
|
||||
|
||||
|
||||
python = sys.executable
|
||||
index_url = os.environ.get("INDEX_URL", "")
|
||||
|
||||
|
||||
def version(package: str) -> Optional[Version]:
|
||||
try:
|
||||
return Version(importlib.import_module(package).__version__)
|
||||
except ModuleNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
def is_installed(package: str) -> bool:
|
||||
try:
|
||||
spec = importlib.util.find_spec(package)
|
||||
except ModuleNotFoundError:
|
||||
return False
|
||||
|
||||
return spec is not None
|
||||
|
||||
|
||||
def run_python(command: str, env: Dict[str, str] = None) -> str:
|
||||
run_kwargs = {
|
||||
"args": f"\"{python}\" {command}",
|
||||
"shell": True,
|
||||
"env": os.environ if env is None else env,
|
||||
"encoding": "utf8",
|
||||
"errors": "ignore",
|
||||
}
|
||||
|
||||
print(run_kwargs["args"])
|
||||
|
||||
result = subprocess.run(**run_kwargs)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error running command: {command}", file=sys.stderr)
|
||||
raise RuntimeError(f"Error running command: {command}")
|
||||
|
||||
return result.stdout or ""
|
||||
|
||||
|
||||
def run_pip(command: str, env: Dict[str, str] = None) -> str:
|
||||
return run_python(f"-m pip {command}", env)
|
||||
@@ -0,0 +1,525 @@
|
||||
import glob
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional, Union, Any, Dict, Tuple, Literal
|
||||
from collections import deque
|
||||
import sys
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.models.optical_flow import raft_small
|
||||
|
||||
from diffusers import LCMScheduler, StableDiffusionPipeline
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import (
|
||||
retrieve_latents,
|
||||
)
|
||||
from .image_utils import postprocess_image, forward_backward_consistency_check
|
||||
from .models.utils import get_nn_latent
|
||||
from .image_filter import SimilarImageFilter
|
||||
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
path_dir = os.path.dirname(CURRENT_DIR)
|
||||
|
||||
def get_instance_path(path):
|
||||
instance_path = os.path.normpath(path)
|
||||
if sys.platform == 'win32':
|
||||
instance_path = instance_path.replace('\\', "/")
|
||||
return instance_path
|
||||
|
||||
class StreamV2V:
|
||||
def __init__(
|
||||
self,
|
||||
pipe: StableDiffusionPipeline,
|
||||
t_index_list: List[int],
|
||||
torch_dtype: torch.dtype = torch.float16,
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
do_add_noise: bool = True,
|
||||
use_denoising_batch: bool = True,
|
||||
frame_buffer_size: int = 1,
|
||||
cfg_type: Literal["none", "full", "self", "initialize"] = "self",
|
||||
) -> None:
|
||||
self.device = pipe.device
|
||||
self.dtype = torch_dtype
|
||||
self.generator = None
|
||||
|
||||
self.height = height
|
||||
self.width = width
|
||||
|
||||
self.latent_height = int(height // pipe.vae_scale_factor)
|
||||
self.latent_width = int(width // pipe.vae_scale_factor)
|
||||
|
||||
self.frame_bff_size = frame_buffer_size
|
||||
self.denoising_steps_num = len(t_index_list)
|
||||
self.cfg_type = cfg_type
|
||||
|
||||
if use_denoising_batch:
|
||||
self.batch_size = self.denoising_steps_num * frame_buffer_size
|
||||
if self.cfg_type == "initialize":
|
||||
self.trt_unet_batch_size = (
|
||||
self.denoising_steps_num + 1
|
||||
) * self.frame_bff_size
|
||||
elif self.cfg_type == "full":
|
||||
self.trt_unet_batch_size = (
|
||||
2 * self.denoising_steps_num * self.frame_bff_size
|
||||
)
|
||||
else:
|
||||
self.trt_unet_batch_size = self.denoising_steps_num * frame_buffer_size
|
||||
else:
|
||||
self.trt_unet_batch_size = self.frame_bff_size
|
||||
self.batch_size = frame_buffer_size
|
||||
|
||||
self.t_list = t_index_list
|
||||
|
||||
self.do_add_noise = do_add_noise
|
||||
self.use_denoising_batch = use_denoising_batch
|
||||
|
||||
self.similar_image_filter = False
|
||||
self.similar_filter = SimilarImageFilter()
|
||||
self.prev_image_tensor = None
|
||||
self.prev_x_t_latent = None
|
||||
self.prev_image_result = None
|
||||
|
||||
self.pipe = pipe
|
||||
self.image_processor = VaeImageProcessor(pipe.vae_scale_factor)
|
||||
|
||||
self.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
|
||||
self.text_encoder = pipe.text_encoder
|
||||
self.unet = pipe.unet
|
||||
self.vae = pipe.vae
|
||||
|
||||
self.cached_x_t_latent = deque(maxlen=4)
|
||||
|
||||
self.inference_time_ema = 0
|
||||
|
||||
def load_lcm_lora(
|
||||
self,
|
||||
pretrained_model_name_or_path_or_dict: Union[
|
||||
str, Dict[str, torch.Tensor]
|
||||
] = "latent-consistency/lcm-lora-sdv1-5",
|
||||
adapter_name: Optional[Any] = 'lcm',
|
||||
**kwargs,
|
||||
) -> None:
|
||||
|
||||
model_path = os.path.join(path_dir,"models")
|
||||
model_lists = os.listdir(model_path)
|
||||
if "lcm-lora-sdv1-5.safetensors" in model_lists: # 加入本地模型支持
|
||||
pretrained_model_name_or_path_or_dict=get_instance_path(os.path.join(model_path,"lcm-lora-sdv1-5.safetensors"))
|
||||
self.pipe.load_lora_weights(
|
||||
pretrained_model_name_or_path_or_dict, adapter_name, **kwargs
|
||||
)
|
||||
|
||||
def load_lora(
|
||||
self,
|
||||
pretrained_lora_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
||||
adapter_name: Optional[Any] = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.pipe.load_lora_weights(
|
||||
pretrained_lora_model_name_or_path_or_dict, adapter_name, **kwargs
|
||||
)
|
||||
|
||||
def fuse_lora(
|
||||
self,
|
||||
fuse_unet: bool = True,
|
||||
fuse_text_encoder: bool = True,
|
||||
lora_scale: float = 1.0,
|
||||
safe_fusing: bool = False,
|
||||
) -> None:
|
||||
self.pipe.fuse_lora(
|
||||
fuse_unet=fuse_unet,
|
||||
fuse_text_encoder=fuse_text_encoder,
|
||||
lora_scale=lora_scale,
|
||||
safe_fusing=safe_fusing,
|
||||
)
|
||||
|
||||
def enable_similar_image_filter(self, threshold: float = 0.98, max_skip_frame: float = 10) -> None:
|
||||
self.similar_image_filter = True
|
||||
self.similar_filter.set_threshold(threshold)
|
||||
self.similar_filter.set_max_skip_frame(max_skip_frame)
|
||||
|
||||
def disable_similar_image_filter(self) -> None:
|
||||
self.similar_image_filter = False
|
||||
|
||||
@torch.no_grad()
|
||||
def prepare(
|
||||
self,
|
||||
prompt: str,
|
||||
negative_prompt: str = "",
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 1.2,
|
||||
delta: float = 1.0,
|
||||
generator: Optional[torch.Generator] = torch.Generator(),
|
||||
seed: int = 2,
|
||||
) -> None:
|
||||
self.generator = generator
|
||||
self.generator.manual_seed(seed)
|
||||
# initialize x_t_latent (it can be any random tensor)
|
||||
if self.denoising_steps_num > 1:
|
||||
self.x_t_latent_buffer = torch.zeros(
|
||||
(
|
||||
(self.denoising_steps_num - 1) * self.frame_bff_size,
|
||||
4,
|
||||
self.latent_height,
|
||||
self.latent_width,
|
||||
),
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
else:
|
||||
self.x_t_latent_buffer = None
|
||||
|
||||
if self.cfg_type == "none":
|
||||
self.guidance_scale = 1.0
|
||||
else:
|
||||
self.guidance_scale = guidance_scale
|
||||
self.delta = delta
|
||||
|
||||
do_classifier_free_guidance = False
|
||||
if self.guidance_scale > 1.0:
|
||||
do_classifier_free_guidance = True
|
||||
|
||||
encoder_output = self.pipe.encode_prompt(
|
||||
prompt=prompt,
|
||||
device=self.device,
|
||||
num_images_per_prompt=1,
|
||||
do_classifier_free_guidance=True,
|
||||
negative_prompt=negative_prompt,
|
||||
)
|
||||
|
||||
self.prompt_embeds = encoder_output[0].repeat(self.batch_size, 1, 1)
|
||||
self.null_prompt_embeds = encoder_output[1]
|
||||
|
||||
if self.use_denoising_batch and self.cfg_type == "full":
|
||||
uncond_prompt_embeds = encoder_output[1].repeat(self.batch_size, 1, 1)
|
||||
elif self.cfg_type == "initialize":
|
||||
uncond_prompt_embeds = encoder_output[1].repeat(self.frame_bff_size, 1, 1)
|
||||
|
||||
if self.guidance_scale > 1.0 and (
|
||||
self.cfg_type == "initialize" or self.cfg_type == "full"
|
||||
):
|
||||
self.prompt_embeds = torch.cat(
|
||||
[uncond_prompt_embeds, self.prompt_embeds], dim=0
|
||||
)
|
||||
|
||||
self.scheduler.set_timesteps(num_inference_steps, self.device)
|
||||
self.timesteps = self.scheduler.timesteps.to(self.device)
|
||||
|
||||
# make sub timesteps list based on the indices in the t_list list and the values in the timesteps list
|
||||
self.sub_timesteps = []
|
||||
for t in self.t_list:
|
||||
self.sub_timesteps.append(self.timesteps[t])
|
||||
|
||||
sub_timesteps_tensor = torch.tensor(
|
||||
self.sub_timesteps, dtype=torch.long, device=self.device
|
||||
)
|
||||
self.sub_timesteps_tensor = torch.repeat_interleave(
|
||||
sub_timesteps_tensor,
|
||||
repeats=self.frame_bff_size if self.use_denoising_batch else 1,
|
||||
dim=0,
|
||||
)
|
||||
|
||||
self.init_noise = torch.randn(
|
||||
(self.batch_size, 4, self.latent_height, self.latent_width),
|
||||
generator=generator,
|
||||
).to(device=self.device, dtype=self.dtype)
|
||||
|
||||
self.randn_noise = self.init_noise[:1].clone()
|
||||
self.warp_noise = self.init_noise[:1].clone()
|
||||
|
||||
self.stock_noise = torch.zeros_like(self.init_noise)
|
||||
|
||||
c_skip_list = []
|
||||
c_out_list = []
|
||||
for timestep in self.sub_timesteps:
|
||||
c_skip, c_out = self.scheduler.get_scalings_for_boundary_condition_discrete(
|
||||
timestep
|
||||
)
|
||||
c_skip_list.append(c_skip)
|
||||
c_out_list.append(c_out)
|
||||
|
||||
self.c_skip = (
|
||||
torch.stack(c_skip_list)
|
||||
.view(len(self.t_list), 1, 1, 1)
|
||||
.to(dtype=self.dtype, device=self.device)
|
||||
)
|
||||
self.c_out = (
|
||||
torch.stack(c_out_list)
|
||||
.view(len(self.t_list), 1, 1, 1)
|
||||
.to(dtype=self.dtype, device=self.device)
|
||||
)
|
||||
|
||||
alpha_prod_t_sqrt_list = []
|
||||
beta_prod_t_sqrt_list = []
|
||||
for timestep in self.sub_timesteps:
|
||||
alpha_prod_t_sqrt = self.scheduler.alphas_cumprod[timestep].sqrt()
|
||||
beta_prod_t_sqrt = (1 - self.scheduler.alphas_cumprod[timestep]).sqrt()
|
||||
alpha_prod_t_sqrt_list.append(alpha_prod_t_sqrt)
|
||||
beta_prod_t_sqrt_list.append(beta_prod_t_sqrt)
|
||||
alpha_prod_t_sqrt = (
|
||||
torch.stack(alpha_prod_t_sqrt_list)
|
||||
.view(len(self.t_list), 1, 1, 1)
|
||||
.to(dtype=self.dtype, device=self.device)
|
||||
)
|
||||
beta_prod_t_sqrt = (
|
||||
torch.stack(beta_prod_t_sqrt_list)
|
||||
.view(len(self.t_list), 1, 1, 1)
|
||||
.to(dtype=self.dtype, device=self.device)
|
||||
)
|
||||
self.alpha_prod_t_sqrt = torch.repeat_interleave(
|
||||
alpha_prod_t_sqrt,
|
||||
repeats=self.frame_bff_size if self.use_denoising_batch else 1,
|
||||
dim=0,
|
||||
)
|
||||
self.beta_prod_t_sqrt = torch.repeat_interleave(
|
||||
beta_prod_t_sqrt,
|
||||
repeats=self.frame_bff_size if self.use_denoising_batch else 1,
|
||||
dim=0,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def update_prompt(self, prompt: str) -> None:
|
||||
encoder_output = self.pipe.encode_prompt(
|
||||
prompt=prompt,
|
||||
device=self.device,
|
||||
num_images_per_prompt=1,
|
||||
do_classifier_free_guidance=False,
|
||||
)
|
||||
self.prompt_embeds = encoder_output[0].repeat(self.batch_size, 1, 1)
|
||||
|
||||
def add_noise(
|
||||
self,
|
||||
original_samples: torch.Tensor,
|
||||
noise: torch.Tensor,
|
||||
t_index: int,
|
||||
) -> torch.Tensor:
|
||||
noisy_samples = (
|
||||
self.alpha_prod_t_sqrt[t_index] * original_samples
|
||||
+ self.beta_prod_t_sqrt[t_index] * noise
|
||||
)
|
||||
return noisy_samples
|
||||
|
||||
def scheduler_step_batch(
|
||||
self,
|
||||
model_pred_batch: torch.Tensor,
|
||||
x_t_latent_batch: torch.Tensor,
|
||||
idx: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
# TODO: use t_list to select beta_prod_t_sqrt
|
||||
if idx is None:
|
||||
F_theta = (
|
||||
x_t_latent_batch - self.beta_prod_t_sqrt * model_pred_batch
|
||||
) / self.alpha_prod_t_sqrt
|
||||
denoised_batch = self.c_out * F_theta + self.c_skip * x_t_latent_batch
|
||||
else:
|
||||
F_theta = (
|
||||
x_t_latent_batch - self.beta_prod_t_sqrt[idx] * model_pred_batch
|
||||
) / self.alpha_prod_t_sqrt[idx]
|
||||
denoised_batch = (
|
||||
self.c_out[idx] * F_theta + self.c_skip[idx] * x_t_latent_batch
|
||||
)
|
||||
|
||||
return denoised_batch
|
||||
|
||||
def unet_step(
|
||||
self,
|
||||
x_t_latent: torch.Tensor,
|
||||
t_list: Union[torch.Tensor, list[int]],
|
||||
idx: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"):
|
||||
x_t_latent_plus_uc = torch.concat([x_t_latent[0:1], x_t_latent], dim=0)
|
||||
t_list = torch.concat([t_list[0:1], t_list], dim=0)
|
||||
elif self.guidance_scale > 1.0 and (self.cfg_type == "full"):
|
||||
x_t_latent_plus_uc = torch.concat([x_t_latent, x_t_latent], dim=0)
|
||||
t_list = torch.concat([t_list, t_list], dim=0)
|
||||
else:
|
||||
x_t_latent_plus_uc = x_t_latent
|
||||
|
||||
model_pred = self.unet(
|
||||
x_t_latent_plus_uc,
|
||||
t_list,
|
||||
encoder_hidden_states=self.prompt_embeds,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"):
|
||||
noise_pred_text = model_pred[1:]
|
||||
self.stock_noise = torch.concat(
|
||||
[model_pred[0:1], self.stock_noise[1:]], dim=0
|
||||
) # ここコメントアウトでself out cfg
|
||||
elif self.guidance_scale > 1.0 and (self.cfg_type == "full"):
|
||||
noise_pred_uncond, noise_pred_text = model_pred.chunk(2)
|
||||
else:
|
||||
noise_pred_text = model_pred
|
||||
if self.guidance_scale > 1.0 and (
|
||||
self.cfg_type == "self" or self.cfg_type == "initialize"
|
||||
):
|
||||
noise_pred_uncond = self.stock_noise * self.delta
|
||||
if self.guidance_scale > 1.0 and self.cfg_type != "none":
|
||||
model_pred = noise_pred_uncond + self.guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
else:
|
||||
model_pred = noise_pred_text
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
if self.use_denoising_batch:
|
||||
denoised_batch = self.scheduler_step_batch(model_pred, x_t_latent, idx)
|
||||
if self.cfg_type == "self" or self.cfg_type == "initialize":
|
||||
scaled_noise = self.beta_prod_t_sqrt * self.stock_noise
|
||||
delta_x = self.scheduler_step_batch(model_pred, scaled_noise, idx)
|
||||
alpha_next = torch.concat(
|
||||
[
|
||||
self.alpha_prod_t_sqrt[1:],
|
||||
torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
delta_x = alpha_next * delta_x
|
||||
beta_next = torch.concat(
|
||||
[
|
||||
self.beta_prod_t_sqrt[1:],
|
||||
torch.ones_like(self.beta_prod_t_sqrt[0:1]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
delta_x = delta_x / beta_next
|
||||
init_noise = torch.concat(
|
||||
[self.init_noise[1:], self.init_noise[0:1]], dim=0
|
||||
)
|
||||
self.stock_noise = init_noise + delta_x
|
||||
|
||||
else:
|
||||
# denoised_batch = self.scheduler.step(model_pred, t_list[0], x_t_latent).denoised
|
||||
denoised_batch = self.scheduler_step_batch(model_pred, x_t_latent, idx)
|
||||
|
||||
return denoised_batch, model_pred
|
||||
|
||||
|
||||
def norm_noise(self, noise):
|
||||
# Compute mean and std of blended_noise
|
||||
mean = noise.mean()
|
||||
std = noise.std()
|
||||
|
||||
# Normalize blended_noise to have mean=0 and std=1
|
||||
normalized_noise = (noise - mean) / std
|
||||
return normalized_noise
|
||||
|
||||
def encode_image(self, image_tensors: torch.Tensor) -> torch.Tensor:
|
||||
image_tensors = image_tensors.to(
|
||||
device=self.device,
|
||||
dtype=self.vae.dtype,
|
||||
)
|
||||
img_latent = retrieve_latents(self.vae.encode(image_tensors), self.generator)
|
||||
img_latent = img_latent * self.vae.config.scaling_factor
|
||||
x_t_latent = self.add_noise(img_latent, self.init_noise[0], 0)
|
||||
return x_t_latent
|
||||
|
||||
def decode_image(self, x_0_pred_out: torch.Tensor) -> torch.Tensor:
|
||||
output_latent = self.vae.decode(
|
||||
x_0_pred_out / self.vae.config.scaling_factor, return_dict=False
|
||||
)[0]
|
||||
return output_latent
|
||||
|
||||
def predict_x0_batch(self, x_t_latent: torch.Tensor) -> torch.Tensor:
|
||||
prev_latent_batch = self.x_t_latent_buffer
|
||||
if self.use_denoising_batch:
|
||||
t_list = self.sub_timesteps_tensor
|
||||
if self.denoising_steps_num > 1:
|
||||
x_t_latent = torch.cat((x_t_latent, prev_latent_batch), dim=0)
|
||||
self.stock_noise = torch.cat(
|
||||
(self.init_noise[0:1], self.stock_noise[:-1]), dim=0
|
||||
)
|
||||
x_0_pred_batch, model_pred = self.unet_step(x_t_latent, t_list)
|
||||
|
||||
if self.denoising_steps_num > 1:
|
||||
x_0_pred_out = x_0_pred_batch[-1].unsqueeze(0)
|
||||
if self.do_add_noise:
|
||||
self.x_t_latent_buffer = (
|
||||
self.alpha_prod_t_sqrt[1:] * x_0_pred_batch[:-1]
|
||||
+ self.beta_prod_t_sqrt[1:] * self.init_noise[1:]
|
||||
)
|
||||
else:
|
||||
self.x_t_latent_buffer = (
|
||||
self.alpha_prod_t_sqrt[1:] * x_0_pred_batch[:-1]
|
||||
)
|
||||
else:
|
||||
x_0_pred_out = x_0_pred_batch
|
||||
self.x_t_latent_buffer = None
|
||||
else:
|
||||
self.init_noise = x_t_latent
|
||||
for idx, t in enumerate(self.sub_timesteps_tensor):
|
||||
t = t.view(
|
||||
1,
|
||||
).repeat(
|
||||
self.frame_bff_size,
|
||||
)
|
||||
x_0_pred, model_pred = self.unet_step(x_t_latent, t, idx)
|
||||
if idx < len(self.sub_timesteps_tensor) - 1:
|
||||
if self.do_add_noise:
|
||||
x_t_latent = self.alpha_prod_t_sqrt[
|
||||
idx + 1
|
||||
] * x_0_pred + self.beta_prod_t_sqrt[
|
||||
idx + 1
|
||||
] * torch.randn_like(
|
||||
x_0_pred, device=self.device, dtype=self.dtype
|
||||
)
|
||||
else:
|
||||
x_t_latent = self.alpha_prod_t_sqrt[idx + 1] * x_0_pred
|
||||
x_0_pred_out = x_0_pred
|
||||
return x_0_pred_out
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self, x: Union[torch.Tensor, PIL.Image.Image, np.ndarray] = None
|
||||
) -> torch.Tensor:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
if x is not None:
|
||||
x = self.image_processor.preprocess(x, self.height, self.width).to(
|
||||
device=self.device, dtype=self.dtype
|
||||
)
|
||||
if self.similar_image_filter:
|
||||
x = self.similar_filter(x)
|
||||
if x is None:
|
||||
time.sleep(self.inference_time_ema)
|
||||
return self.prev_image_result
|
||||
x_t_latent = self.encode_image(x)
|
||||
else:
|
||||
# TODO: check the dimension of x_t_latent
|
||||
x_t_latent = torch.randn((1, 4, self.latent_height, self.latent_width)).to(
|
||||
device=self.device, dtype=self.dtype
|
||||
)
|
||||
x_0_pred_out = self.predict_x0_batch(x_t_latent)
|
||||
x_output = self.decode_image(x_0_pred_out).detach().clone()
|
||||
|
||||
self.prev_image_result = x_output
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
inference_time = start.elapsed_time(end) / 1000
|
||||
self.inference_time_ema = 0.9 * self.inference_time_ema + 0.1 * inference_time
|
||||
return x_output
|
||||
|
||||
@torch.no_grad()
|
||||
def txt2img(self, batch_size: int = 1) -> torch.Tensor:
|
||||
x_0_pred_out = self.predict_x0_batch(self.init_noise[0:1])
|
||||
x_output = self.decode_image(x_0_pred_out).detach().clone()
|
||||
return x_output
|
||||
|
||||
def txt2img_sd_turbo(self, batch_size: int = 1) -> torch.Tensor:
|
||||
x_t_latent = self.init_noise[0:1]
|
||||
model_pred = self.unet(
|
||||
x_t_latent,
|
||||
self.sub_timesteps_tensor,
|
||||
encoder_hidden_states=self.prompt_embeds,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
x_0_pred_out = (
|
||||
x_t_latent - self.beta_prod_t_sqrt * model_pred
|
||||
) / self.alpha_prod_t_sqrt
|
||||
return self.decode_image(x_0_pred_out)
|
||||
Reference in New Issue
Block a user