419 lines
14 KiB
Python
419 lines
14 KiB
Python
import math
|
|
import sys
|
|
from typing import Callable, List, Optional, Union
|
|
|
|
import torch
|
|
from einops import rearrange, repeat
|
|
from torch import Tensor
|
|
|
|
from ..models.model import Flux
|
|
from ..modules.conditioner import HFEmbedder
|
|
from ..modules.image_embedders import ReduxImageEncoder
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Progress bar
|
|
# -------------------------------------------------------------------------
|
|
import time
|
|
|
|
TGT_PREFIX = "[TARGET-SCENE]"
|
|
|
|
|
|
def print_progress_bar(iteration, total, prefix='', suffix='', length=30, fill='█'):
|
|
"""
|
|
Simple progress bar for console output, with elapsed and estimated remaining time.
|
|
|
|
Args:
|
|
iteration: Current iteration (Int)
|
|
total: Total iterations (Int)
|
|
prefix: Prefix string (Str)
|
|
suffix: Suffix string (Str)
|
|
length: Bar length (Int)
|
|
fill: Bar fill character (Str)
|
|
"""
|
|
# Static variable to store start time
|
|
if not hasattr(print_progress_bar, "_start_time") or iteration == 0:
|
|
print_progress_bar._start_time = time.time()
|
|
|
|
percent = f"{100 * (iteration / float(total)):.1f}%"
|
|
filled_length = int(length * iteration // total)
|
|
bar = fill * filled_length + '-' * (length - filled_length)
|
|
|
|
elapsed = time.time() - print_progress_bar._start_time
|
|
elapsed_str = time.strftime("%H:%M:%S", time.gmtime(elapsed))
|
|
|
|
if iteration > 0:
|
|
avg_time_per_iter = elapsed / iteration
|
|
remaining = avg_time_per_iter * (total - iteration)
|
|
else:
|
|
remaining = 0
|
|
remaining_str = time.strftime("%H:%M:%S", time.gmtime(remaining))
|
|
|
|
time_info = f"Elapsed: {elapsed_str} | ETA: {remaining_str}"
|
|
|
|
sys.stdout.write(f'\r{prefix} |{bar}| {percent} {suffix} {time_info}')
|
|
sys.stdout.flush()
|
|
if iteration == total:
|
|
sys.stdout.write('\n')
|
|
sys.stdout.flush()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 1) sampling func
|
|
# -------------------------------------------------------------------------
|
|
def unpack(x: Tensor, height: int, width: int) -> Tensor:
|
|
return rearrange(
|
|
x,
|
|
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
|
h=math.ceil(height / 16),
|
|
w=math.ceil(width / 16),
|
|
ph=2,
|
|
pw=2,
|
|
)
|
|
|
|
|
|
def time_shift(mu: float, sigma: float, t: Tensor):
|
|
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
|
|
|
|
|
def get_lin_function(
|
|
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
|
):
|
|
m = (y2 - y1) / (x2 - x1)
|
|
b = y1 - m * x1
|
|
return lambda x: m * x + b
|
|
|
|
|
|
def get_schedule(
|
|
num_steps: int,
|
|
image_seq_len: int,
|
|
base_shift: float = 0.5,
|
|
max_shift: float = 1.15,
|
|
shift: bool = True,
|
|
):
|
|
# extra step for zero
|
|
timesteps = torch.linspace(1, 0, num_steps + 1)
|
|
|
|
# shifting the schedule to favor high timesteps for higher signal images
|
|
if shift:
|
|
# eastimate mu based on linear estimation between two points
|
|
mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
|
|
timesteps = time_shift(mu, 1.0, timesteps)
|
|
|
|
return timesteps.tolist()
|
|
|
|
|
|
def get_noise(
|
|
num_samples: int,
|
|
height: int,
|
|
width: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
seed: int,
|
|
):
|
|
|
|
noise = torch.cat(
|
|
[torch.randn(
|
|
1,
|
|
16,
|
|
# allow for packing
|
|
2 * math.ceil(height / 16),
|
|
2 * math.ceil(width / 16),
|
|
device=device,
|
|
dtype=dtype,
|
|
generator=torch.Generator(device=device).manual_seed(seed+i),
|
|
)
|
|
for i in range(num_samples)
|
|
],
|
|
dim=0
|
|
)
|
|
return noise
|
|
|
|
|
|
|
|
# -------------------------------------------------------------------------
|
|
# prepare input func
|
|
# -------------------------------------------------------------------------
|
|
def _get_batch_size_and_prompt(prompt, img_shape):
|
|
"""
|
|
Helper to determine batch size and prompt list.
|
|
"""
|
|
bs, c, h, w = img_shape
|
|
is_prompt_none = prompt is None
|
|
return bs, prompt, is_prompt_none, h, w
|
|
|
|
def _make_img_ids(bs, h, w, device=None, dtype=None):
|
|
"""
|
|
Helper to create image ids tensor.
|
|
"""
|
|
img_ids = torch.zeros(h // 2, w // 2, 3, device=device, dtype=dtype)
|
|
img_ids[..., 1] = torch.arange(h // 2)[:, None]
|
|
img_ids[..., 2] = torch.arange(w // 2)[None, :]
|
|
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
|
return img_ids
|
|
|
|
|
|
def prepare(
|
|
t5: HFEmbedder,
|
|
clip: HFEmbedder,
|
|
img: Tensor,
|
|
prompt: Union[str, List[str], None],
|
|
num_images_per_prompt: int = 1,
|
|
):
|
|
"""
|
|
Prepare the regular input for the Diffusion Transformer.
|
|
"""
|
|
img_bs, prompt, is_prompt_none, h, w = _get_batch_size_and_prompt(prompt, img.shape)
|
|
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
|
img_ids = _make_img_ids(img_bs, h, w, device=img.device, dtype=img.dtype)
|
|
|
|
if isinstance(prompt, str):
|
|
prompt = [prompt]
|
|
|
|
txt_bs = len(prompt)
|
|
if not is_prompt_none:
|
|
prompt = [TGT_PREFIX + p for p in prompt]
|
|
txt = t5(prompt)
|
|
txt_ids = torch.zeros(txt_bs, txt.shape[1], 3, device=img.device, dtype=img.dtype)
|
|
txt_vec = clip(prompt)
|
|
else:
|
|
|
|
|
|
txt = torch.zeros(txt_bs, 512, 4096, device=img.device, dtype=img.dtype)
|
|
txt_ids = torch.zeros(txt_bs, 512, 3, device=img.device, dtype=img.dtype)
|
|
txt_vec = torch.zeros(txt_bs, 768, device=img.device, dtype=img.dtype)
|
|
|
|
|
|
if num_images_per_prompt > 1:
|
|
txt = txt.repeat_interleave(num_images_per_prompt, dim=0)
|
|
txt_ids = txt_ids.repeat_interleave(num_images_per_prompt, dim=0)
|
|
txt_vec = txt_vec.repeat_interleave(num_images_per_prompt, dim=0)
|
|
|
|
|
|
return {
|
|
"img": img.to(img.device),
|
|
"img_ids": img_ids.to(img.device),
|
|
"txt": txt.to(img.device),
|
|
"txt_ids": txt_ids.to(img.device),
|
|
"txt_vec": txt_vec.to(img.device),
|
|
}
|
|
|
|
|
|
def prepare_with_redux(
|
|
t5: HFEmbedder,
|
|
clip: HFEmbedder,
|
|
image_encoder: ReduxImageEncoder,
|
|
img: Tensor,
|
|
img_ip: Tensor,
|
|
prompt: Union[str, List[str], None],
|
|
num_images_per_prompt: int = 1,
|
|
):
|
|
img_bs, prompt, is_prompt_none, h, w = _get_batch_size_and_prompt(prompt, img.shape)
|
|
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
|
img_ids = _make_img_ids(img_bs, h, w, device=img.device, dtype=img.dtype)
|
|
|
|
if isinstance(prompt, str):
|
|
prompt = [prompt]
|
|
txt_bs = len(prompt)
|
|
|
|
if not is_prompt_none:
|
|
prompt = [TGT_PREFIX + p for p in prompt]
|
|
txt = torch.cat((t5(prompt), image_encoder(img_ip)), dim=1)
|
|
txt_ids = torch.zeros(txt_bs, txt.shape[1], 3, device=img.device, dtype=img.dtype)
|
|
txt_vec = clip(prompt)
|
|
else:
|
|
txt = torch.zeros(txt_bs, 512, 4096, device=img.device, dtype=img.dtype)
|
|
txt_ids = torch.zeros(txt_bs, 512, 3, device=img.device, dtype=img.dtype)
|
|
txt_vec = torch.zeros(txt_bs, 768, device=img.device, dtype=img.dtype)
|
|
|
|
if num_images_per_prompt > 1:
|
|
txt = txt.repeat_interleave(num_images_per_prompt, dim=0)
|
|
txt_ids = txt_ids.repeat_interleave(num_images_per_prompt, dim=0)
|
|
txt_vec = txt_vec.repeat_interleave(num_images_per_prompt, dim=0)
|
|
|
|
return {
|
|
"img": img.to(img.device),
|
|
"img_ids": img_ids.to(img.device),
|
|
"txt": txt.to(img.device),
|
|
"txt_ids": txt_ids.to(img.device),
|
|
"txt_vec": txt_vec.to(img.device),
|
|
}
|
|
|
|
|
|
|
|
def prepare_image_cond(
|
|
ae,
|
|
img_ref,
|
|
img_target,
|
|
mask_target,
|
|
dtype,
|
|
device,
|
|
num_images_per_prompt: int = 1,
|
|
):
|
|
batch_size, _, _, _ = img_target.shape
|
|
|
|
# Apply mask to target image
|
|
mask_targeted_img = img_target * mask_target
|
|
|
|
if mask_target.shape[1] == 3:
|
|
mask_target = mask_target[:, 0 : 1, :, :]
|
|
|
|
with torch.no_grad():
|
|
autoencoder_dtype = next(ae.parameters()).dtype
|
|
# Encode masked target image to latent space
|
|
mask_targeted_latent = ae.encode(mask_targeted_img.to(autoencoder_dtype)).to(dtype)
|
|
# Encode reference image to latent space
|
|
reference_latent = ae.encode(img_ref.to(autoencoder_dtype)).to(dtype)
|
|
|
|
# Repeat reference latent if batch size > 1
|
|
if reference_latent.shape[0] == 1 and batch_size > 1:
|
|
reference_latent = repeat(reference_latent, "1 ... -> bs ...", bs=batch_size)
|
|
|
|
# Concatenate reference and target latents
|
|
latent_concat = torch.cat((reference_latent, mask_targeted_latent), dim=-1)
|
|
# Pack latents into 2x2 patches
|
|
latent_packed = rearrange(latent_concat, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
|
|
|
# Create reference mask (all ones)
|
|
reference_mask = torch.ones_like(img_ref)
|
|
if reference_mask.shape[1] == 3:
|
|
reference_mask = reference_mask[:, 0 : 1, :, :]
|
|
|
|
# Concatenate reference and target masks
|
|
mask_concat = torch.cat((reference_mask, mask_target), dim=-1)
|
|
# Pack masks into 16x16 patches for image conditioning
|
|
mask_16x16 = rearrange(mask_concat, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=16, pw=16)
|
|
|
|
# Interpolate masks to latent space dimensions
|
|
mask_latent = torch.nn.functional.interpolate(mask_concat, size=(latent_concat.shape[2] // 2, latent_concat.shape[3] // 2), mode='nearest')
|
|
|
|
# Pack interpolated masks into 1x1 patches for mask conditioning
|
|
mask_cond = rearrange(mask_latent, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=1, pw=1)
|
|
# Combine packed latents and masks for image conditioning
|
|
img_cond = torch.cat((latent_packed, mask_16x16), dim=-1)
|
|
|
|
if num_images_per_prompt > 1:
|
|
img_cond = img_cond.repeat_interleave(num_images_per_prompt, dim=0)
|
|
mask_cond = mask_cond.repeat_interleave(num_images_per_prompt, dim=0)
|
|
latent_packed = latent_packed.repeat_interleave(num_images_per_prompt, dim=0)
|
|
|
|
return {
|
|
"img_cond": img_cond.to(device).to(dtype),
|
|
"mask_cond": mask_cond.to(device).to(dtype),
|
|
"img_latent": latent_packed.to(device).to(dtype),
|
|
}
|
|
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 2) denoise func
|
|
# -------------------------------------------------------------------------
|
|
def is_even_step(step: int) -> bool:
|
|
"""Check if the current step is odd."""
|
|
return (step % 2 == 0)
|
|
|
|
def denoise(
|
|
model,
|
|
img,
|
|
img_ids,
|
|
txt,
|
|
txt_ids,
|
|
txt_vec,
|
|
timesteps,
|
|
guidance: float = 4.0,
|
|
img_cond: Tensor = None,
|
|
mask_cond: Tensor = None,
|
|
img_latent: Tensor = None,
|
|
cond_w_regions: Optional[Union[List[int], int]] = None,
|
|
mask_type_ids: Optional[Union[Tensor, int]] = None,
|
|
height: int = 1024,
|
|
width: int = 1024,
|
|
use_background_preservation: bool = False,
|
|
use_progressive_background_preservation: bool = True,
|
|
background_blend_threshold: float = 0.8,
|
|
true_gs: float = 1,
|
|
timestep_to_start_cfg: int = 0,
|
|
neg_txt: Tensor = None,
|
|
neg_txt_ids: Tensor = None,
|
|
neg_txt_vec: Tensor = None,
|
|
show_progress: bool = False,
|
|
use_flash_attention: bool = False,
|
|
gradio_progress=None,
|
|
**kwargs,
|
|
):
|
|
do_true_cfg = true_gs > 1 and neg_txt is not None
|
|
|
|
guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
|
|
v_gt = img - img_latent
|
|
num_steps = len(timesteps[:-1])
|
|
|
|
update_func = kwargs.get('update_func', lambda *args, **kwargs: None)
|
|
|
|
for step, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
|
|
update_func()
|
|
if show_progress:
|
|
print_progress_bar(step, num_steps, prefix='Denoising:', suffix=f'Step {step+1}/{num_steps}')
|
|
|
|
# Update Gradio progress if available
|
|
if gradio_progress is not None:
|
|
# Map denoise progress to 0.2-0.8 range (since 0.0-0.2 is preprocessing, 0.8-1.0 is postprocessing)
|
|
progress_value = (step / num_steps)
|
|
gradio_progress(progress_value, desc=f"Denoising step {step+1}/{num_steps}")
|
|
|
|
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
|
model_dtype = list(model.parameters())[0].dtype
|
|
|
|
pred = model(
|
|
img=torch.cat((img.to(model_dtype), img_cond.to(model_dtype)), dim=-1) if img_cond is not None else img.to(model_dtype),
|
|
img_ids=img_ids.to(model_dtype),
|
|
txt=txt.to(model_dtype),
|
|
txt_ids=txt_ids.to(model_dtype),
|
|
txt_vec=txt_vec.to(model_dtype),
|
|
timesteps=t_vec.to(model_dtype),
|
|
guidance=guidance_vec.to(model_dtype),
|
|
cond_w_regions=cond_w_regions,
|
|
mask_type_ids=mask_type_ids,
|
|
height=height,
|
|
width=width,
|
|
use_flash_attention=use_flash_attention,
|
|
)
|
|
|
|
if do_true_cfg and step >= timestep_to_start_cfg:
|
|
neg_perd = model(
|
|
img=torch.cat((img.to(model_dtype), img_cond.to(model_dtype)), dim=-1) if img_cond is not None else img.to(model_dtype),
|
|
img_ids=img_ids.to(model_dtype),
|
|
txt=neg_txt.to(model_dtype),
|
|
txt_ids=neg_txt_ids.to(model_dtype),
|
|
txt_vec=neg_txt_vec.to(model_dtype),
|
|
timesteps=t_vec.to(model_dtype),
|
|
guidance=guidance_vec.to(model_dtype),
|
|
cond_w_regions=cond_w_regions,
|
|
mask_type_ids=mask_type_ids,
|
|
height=height,
|
|
width=width,
|
|
use_flash_attention=use_flash_attention,
|
|
)
|
|
pred = neg_perd + true_gs * (pred - neg_perd)
|
|
|
|
if use_background_preservation:
|
|
is_early_phase = step <= num_steps * background_blend_threshold
|
|
|
|
if is_early_phase:
|
|
if use_progressive_background_preservation:
|
|
if is_even_step(step):
|
|
# Apply mask blending on odd steps in early phase
|
|
masked_latent = pred * (1 - mask_cond) + v_gt * mask_cond
|
|
else:
|
|
# Use prediction directly for even steps or late phase
|
|
masked_latent = pred
|
|
else:
|
|
masked_latent = pred * (1 - mask_cond) + v_gt * mask_cond
|
|
else:
|
|
# Use prediction directly for even steps or late phase
|
|
masked_latent = pred
|
|
|
|
img = img + (t_prev - t_curr) * masked_latent
|
|
else:
|
|
img = img + (t_prev - t_curr) * pred
|
|
|
|
if show_progress:
|
|
print_progress_bar(num_steps, num_steps, prefix='Denoising:', suffix='Complete')
|
|
|
|
return img |