|
|
from contextlib import nullcontext |
|
|
from typing import List, Optional, Tuple, Union |
|
|
|
|
|
import torch |
|
|
from einops import rearrange |
|
|
from PIL import Image |
|
|
from tqdm.auto import tqdm |
|
|
|
|
|
from diffusers import DiffusionPipeline |
|
|
from diffusers.pipelines.pipeline_utils import ImagePipelineOutput |
|
|
|
|
|
from .constants import SUPPORTED_IMAGE_SIZES |
|
|
|
|
|
|
|
|
PromptType = Union[str, List[str]] |
|
|
|
|
|
|
|
|
def _get_pkv_seq_len(past_key_values) -> int: |
|
|
"""Get cached sequence length from past_key_values (supports tuple and DynamicCache).""" |
|
|
if hasattr(past_key_values, "get_seq_length"): |
|
|
return past_key_values.get_seq_length() |
|
|
return past_key_values[0][0].shape[2] |
|
|
|
|
|
|
|
|
class BitDanceDiffusionPipeline(DiffusionPipeline): |
|
|
model_cpu_offload_seq = "text_encoder->projector->diffusion_head->autoencoder" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
tokenizer, |
|
|
text_encoder, |
|
|
autoencoder, |
|
|
diffusion_head, |
|
|
projector, |
|
|
supported_image_sizes: Optional[List[List[int]]] = None, |
|
|
dtype: Optional[torch.dtype] = None, |
|
|
) -> None: |
|
|
super().__init__() |
|
|
self.register_modules( |
|
|
tokenizer=tokenizer, |
|
|
text_encoder=text_encoder, |
|
|
autoencoder=autoencoder, |
|
|
diffusion_head=diffusion_head, |
|
|
projector=projector, |
|
|
) |
|
|
|
|
|
image_sizes = supported_image_sizes or SUPPORTED_IMAGE_SIZES |
|
|
self.register_to_config(supported_image_sizes=[list(size) for size in image_sizes]) |
|
|
|
|
|
self.hidden_size = self.text_encoder.config.hidden_size |
|
|
self.vae_patch_size = self.autoencoder.patch_size |
|
|
self.parallel_num = int(self.diffusion_head.config.parallel_num) |
|
|
self.ps = int(self.parallel_num**0.5) |
|
|
if self.ps * self.ps != self.parallel_num: |
|
|
raise ValueError( |
|
|
f"parallel_num must be a perfect square (got {self.parallel_num})." |
|
|
) |
|
|
|
|
|
self._build_pos_embed() |
|
|
|
|
|
@property |
|
|
def supported_image_sizes(self) -> List[List[int]]: |
|
|
return [list(size) for size in self.config.supported_image_sizes] |
|
|
|
|
|
def _execution_device_fallback(self) -> torch.device: |
|
|
if getattr(self, "_execution_device", None) is not None: |
|
|
return self._execution_device |
|
|
return next(self.text_encoder.parameters()).device |
|
|
|
|
|
def _build_pos_embed(self) -> None: |
|
|
max_resolution = max(max(size) for size in self.supported_image_sizes) |
|
|
max_len = max_resolution // self.vae_patch_size |
|
|
pos_embed_1d = self._get_1d_sincos_pos_embed(self.hidden_size // 2, max_len) |
|
|
self.pos_embed_1d = pos_embed_1d |
|
|
|
|
|
@staticmethod |
|
|
def _get_1d_sincos_pos_embed(dim: int, max_len: int, pe_interpolation: float = 1.0) -> torch.Tensor: |
|
|
if dim % 2 != 0: |
|
|
raise ValueError(f"dim must be even, got {dim}") |
|
|
omega = torch.arange(dim // 2, dtype=torch.float32) |
|
|
omega /= dim / 2.0 |
|
|
omega = 1.0 / 10000**omega |
|
|
pos = torch.arange(max_len, dtype=torch.float32) / pe_interpolation |
|
|
out = torch.einsum("m,d->md", pos, omega) |
|
|
emb_sin = torch.sin(out) |
|
|
emb_cos = torch.cos(out) |
|
|
return torch.cat([emb_sin, emb_cos], dim=1) |
|
|
|
|
|
def _get_2d_embed(self, h: int, w: int, ps: int = 1) -> torch.Tensor: |
|
|
emb_v = self.pos_embed_1d[:h] |
|
|
emb_h = self.pos_embed_1d[:w] |
|
|
grid_v = emb_v.view(h, 1, self.hidden_size // 2).repeat(1, w, 1) |
|
|
grid_h = emb_h.view(1, w, self.hidden_size // 2).repeat(h, 1, 1) |
|
|
pos_embed = torch.cat([grid_h, grid_v], dim=-1) |
|
|
return rearrange(pos_embed, "(h p1) (w p2) c -> (h w p1 p2) c", p1=ps, p2=ps) |
|
|
|
|
|
def _encode_prompt_to_embeds( |
|
|
self, |
|
|
prompt: str, |
|
|
image_size: Tuple[int, int], |
|
|
num_images_per_prompt: int, |
|
|
guidance_scale: float, |
|
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: |
|
|
device = self._execution_device_fallback() |
|
|
model = self.text_encoder.model |
|
|
tokenizer = self.tokenizer |
|
|
|
|
|
cond_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" |
|
|
uncond_prompt = "<|im_start|>assistant\n" |
|
|
|
|
|
cond_ids = torch.tensor(tokenizer.encode(cond_prompt), device=device, dtype=torch.long) |
|
|
cond_emb = model.embed_tokens(cond_ids) |
|
|
uncond_emb = None |
|
|
if guidance_scale > 1.0: |
|
|
uncond_ids = torch.tensor(tokenizer.encode(uncond_prompt), device=device, dtype=torch.long) |
|
|
uncond_emb = model.embed_tokens(uncond_ids) |
|
|
|
|
|
image_h, image_w = image_size |
|
|
img_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>") |
|
|
res_h_token_id = tokenizer.convert_tokens_to_ids(f"<|res_{image_h // self.vae_patch_size}|>") |
|
|
res_w_token_id = tokenizer.convert_tokens_to_ids(f"<|res_{image_w // self.vae_patch_size}|>") |
|
|
img_start_emb = model.embed_tokens(torch.tensor([img_start_id, res_h_token_id, res_w_token_id], device=device)) |
|
|
|
|
|
for i in range(1, self.parallel_num): |
|
|
query_token_id = tokenizer.convert_tokens_to_ids(f"<|query_{i}|>") |
|
|
query_token = torch.tensor([query_token_id], device=device, dtype=torch.long) |
|
|
query_embed = model.embed_tokens(query_token) |
|
|
img_start_emb = torch.cat([img_start_emb, query_embed], dim=0) |
|
|
|
|
|
input_embeds_cond = torch.cat([cond_emb, img_start_emb], dim=0).unsqueeze(0).repeat(num_images_per_prompt, 1, 1) |
|
|
input_embeds_uncond = None |
|
|
if guidance_scale > 1.0 and uncond_emb is not None: |
|
|
input_embeds_uncond = torch.cat([uncond_emb, img_start_emb], dim=0).unsqueeze(0).repeat(num_images_per_prompt, 1, 1) |
|
|
return input_embeds_cond, input_embeds_uncond, img_start_emb |
|
|
|
|
|
def _decode_tokens_to_image(self, image_latents: torch.Tensor, image_size: Tuple[int, int], ps: int = 1) -> torch.Tensor: |
|
|
h, w = image_size |
|
|
image_latents = rearrange(image_latents, "b (h w p1 p2) c -> b c (h p1) (w p2)", h=h // ps, w=w // ps, p1=ps, p2=ps) |
|
|
return self.autoencoder.decode(image_latents) |
|
|
|
|
|
@torch.no_grad() |
|
|
def _generate_single_prompt( |
|
|
self, |
|
|
prompt: str, |
|
|
height: int, |
|
|
width: int, |
|
|
num_inference_steps: int, |
|
|
guidance_scale: float, |
|
|
num_images_per_prompt: int, |
|
|
generator: Optional[torch.Generator], |
|
|
show_progress_bar: bool, |
|
|
) -> torch.Tensor: |
|
|
image_size = (height, width) |
|
|
if list(image_size) not in self.supported_image_sizes: |
|
|
raise ValueError( |
|
|
f"image_size {list(image_size)} is not supported. " |
|
|
f"Please choose from {self.supported_image_sizes}" |
|
|
) |
|
|
|
|
|
h, w = height // self.vae_patch_size, width // self.vae_patch_size |
|
|
max_length = h * w |
|
|
step_width = self.parallel_num |
|
|
if max_length % step_width != 0: |
|
|
raise ValueError( |
|
|
f"max_length ({max_length}) must be divisible by parallel_num ({step_width})." |
|
|
) |
|
|
num_steps = max_length // step_width |
|
|
|
|
|
device = self._execution_device_fallback() |
|
|
model = self.text_encoder.model |
|
|
dtype = next(self.text_encoder.parameters()).dtype |
|
|
|
|
|
input_embeds_cond, input_embeds_uncond, _ = self._encode_prompt_to_embeds( |
|
|
prompt=prompt, |
|
|
image_size=image_size, |
|
|
num_images_per_prompt=num_images_per_prompt, |
|
|
guidance_scale=guidance_scale, |
|
|
) |
|
|
pos_embed_for_diff = self._get_2d_embed(h, w, ps=self.ps).unsqueeze(0).to(device=device, dtype=dtype) |
|
|
|
|
|
autocast_ctx = ( |
|
|
torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16) |
|
|
if device.type == "cuda" |
|
|
else nullcontext() |
|
|
) |
|
|
|
|
|
with autocast_ctx: |
|
|
outputs_c = model(inputs_embeds=input_embeds_cond[:, :-step_width, :], use_cache=True) |
|
|
pkv_c = outputs_c.past_key_values |
|
|
|
|
|
bi_attn_mask = torch.ones( |
|
|
(input_embeds_cond.shape[0], 1, step_width, step_width + _get_pkv_seq_len(pkv_c)), |
|
|
dtype=torch.bool, |
|
|
device=device, |
|
|
) |
|
|
outputs_c = model( |
|
|
inputs_embeds=input_embeds_cond[:, -step_width:, :], |
|
|
past_key_values=pkv_c, |
|
|
use_cache=True, |
|
|
attention_mask=bi_attn_mask, |
|
|
) |
|
|
pkv_c = outputs_c.past_key_values |
|
|
hidden_c = outputs_c.last_hidden_state[:, -step_width:] |
|
|
|
|
|
hidden_u = None |
|
|
pkv_u = None |
|
|
if guidance_scale > 1.0 and input_embeds_uncond is not None: |
|
|
outputs_u = model(inputs_embeds=input_embeds_uncond[:, :-step_width, :], use_cache=True) |
|
|
pkv_u = outputs_u.past_key_values |
|
|
bi_attn_mask_u = torch.ones( |
|
|
(input_embeds_uncond.shape[0], 1, step_width, step_width + _get_pkv_seq_len(pkv_u)), |
|
|
dtype=torch.bool, |
|
|
device=device, |
|
|
) |
|
|
outputs_u = model( |
|
|
inputs_embeds=input_embeds_uncond[:, -step_width:, :], |
|
|
past_key_values=pkv_u, |
|
|
use_cache=True, |
|
|
attention_mask=bi_attn_mask_u, |
|
|
) |
|
|
pkv_u = outputs_u.past_key_values |
|
|
hidden_u = outputs_u.last_hidden_state[:, -step_width:] |
|
|
|
|
|
out_tokens = [] |
|
|
step_iter = range(num_steps) |
|
|
if show_progress_bar: |
|
|
step_iter = tqdm(step_iter, total=num_steps, desc="Decoding steps") |
|
|
|
|
|
for step in step_iter: |
|
|
if guidance_scale > 1.0 and hidden_u is not None: |
|
|
h_fused = torch.cat([hidden_c, hidden_u], dim=0) |
|
|
else: |
|
|
h_fused = hidden_c |
|
|
|
|
|
pos_slice = pos_embed_for_diff[:, step * step_width : (step + 1) * step_width, :] |
|
|
h_fused = h_fused + pos_slice |
|
|
pred_latents = self.diffusion_head.sample( |
|
|
h_fused, |
|
|
num_sampling_steps=num_inference_steps, |
|
|
cfg=guidance_scale, |
|
|
generator=generator, |
|
|
) |
|
|
curr_tokens = torch.sign(pred_latents) |
|
|
curr_embeds = self.projector(curr_tokens) |
|
|
out_tokens.append(curr_tokens[:num_images_per_prompt]) |
|
|
|
|
|
model_input = curr_embeds + pos_slice |
|
|
bi_attn_mask = torch.ones( |
|
|
(model_input.shape[0], 1, model_input.shape[1], model_input.shape[1] + _get_pkv_seq_len(pkv_c)), |
|
|
dtype=torch.bool, |
|
|
device=device, |
|
|
) |
|
|
outputs_c = model( |
|
|
inputs_embeds=model_input[:num_images_per_prompt], |
|
|
past_key_values=pkv_c, |
|
|
use_cache=True, |
|
|
attention_mask=bi_attn_mask[:num_images_per_prompt], |
|
|
) |
|
|
pkv_c = outputs_c.past_key_values |
|
|
hidden_c = outputs_c.last_hidden_state[:, -step_width:] |
|
|
|
|
|
if guidance_scale > 1.0 and hidden_u is not None and pkv_u is not None: |
|
|
bi_attn_mask_u = torch.ones( |
|
|
(model_input.shape[0], 1, model_input.shape[1], model_input.shape[1] + _get_pkv_seq_len(pkv_u)), |
|
|
dtype=torch.bool, |
|
|
device=device, |
|
|
) |
|
|
outputs_u = model( |
|
|
inputs_embeds=model_input[num_images_per_prompt:], |
|
|
past_key_values=pkv_u, |
|
|
use_cache=True, |
|
|
attention_mask=bi_attn_mask_u[num_images_per_prompt:], |
|
|
) |
|
|
pkv_u = outputs_u.past_key_values |
|
|
hidden_u = outputs_u.last_hidden_state[:, -step_width:] |
|
|
|
|
|
full_output = torch.cat(out_tokens, dim=1) |
|
|
return self._decode_tokens_to_image(full_output, image_size=(h, w), ps=self.ps) |
|
|
|
|
|
@torch.no_grad() |
|
|
def __call__( |
|
|
self, |
|
|
prompt: PromptType, |
|
|
height: int = 1024, |
|
|
width: int = 1024, |
|
|
num_inference_steps: int = 50, |
|
|
guidance_scale: float = 7.5, |
|
|
num_images_per_prompt: int = 1, |
|
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, |
|
|
output_type: str = "pil", |
|
|
return_dict: bool = True, |
|
|
show_progress_bar: bool = False, |
|
|
) -> Union[ImagePipelineOutput, Tuple]: |
|
|
prompts = [prompt] if isinstance(prompt, str) else list(prompt) |
|
|
if len(prompts) == 0: |
|
|
raise ValueError("prompt must be a non-empty string or list of strings.") |
|
|
|
|
|
if isinstance(generator, list) and len(generator) != len(prompts): |
|
|
raise ValueError("When passing a list of generators, its length must equal len(prompt).") |
|
|
|
|
|
image_tensors = [] |
|
|
for i, prompt_text in enumerate(prompts): |
|
|
prompt_generator = generator[i] if isinstance(generator, list) else generator |
|
|
images = self._generate_single_prompt( |
|
|
prompt=prompt_text, |
|
|
height=height, |
|
|
width=width, |
|
|
num_inference_steps=num_inference_steps, |
|
|
guidance_scale=guidance_scale, |
|
|
num_images_per_prompt=num_images_per_prompt, |
|
|
generator=prompt_generator, |
|
|
show_progress_bar=show_progress_bar, |
|
|
) |
|
|
image_tensors.append(images) |
|
|
|
|
|
images_pt = torch.cat(image_tensors, dim=0) |
|
|
images_pt_01 = torch.clamp((images_pt + 1.0) / 2.0, 0.0, 1.0) |
|
|
|
|
|
if output_type == "pt": |
|
|
output_images = images_pt_01 |
|
|
elif output_type == "np": |
|
|
output_images = images_pt_01.permute(0, 2, 3, 1).float().cpu().numpy() |
|
|
elif output_type == "pil": |
|
|
images_uint8 = ( |
|
|
torch.clamp(127.5 * images_pt + 128.0, 0, 255) |
|
|
.permute(0, 2, 3, 1) |
|
|
.to("cpu", dtype=torch.uint8) |
|
|
.numpy() |
|
|
) |
|
|
output_images = [Image.fromarray(image) for image in images_uint8] |
|
|
else: |
|
|
raise ValueError(f"Unsupported output_type={output_type}. Expected 'pil', 'np', or 'pt'.") |
|
|
|
|
|
if not return_dict: |
|
|
return (output_images,) |
|
|
return ImagePipelineOutput(images=output_images) |
|
|
|
|
|
@torch.no_grad() |
|
|
def generate( |
|
|
self, |
|
|
prompt: str, |
|
|
height: int = 1024, |
|
|
width: int = 1024, |
|
|
num_sampling_steps: int = 50, |
|
|
guidance_scale: float = 7.5, |
|
|
num_images: int = 1, |
|
|
seed: Optional[int] = None, |
|
|
) -> List[Image.Image]: |
|
|
generator = None |
|
|
if seed is not None: |
|
|
device = self._execution_device_fallback() |
|
|
generator_device = "cuda" if device.type == "cuda" else "cpu" |
|
|
generator = torch.Generator(device=generator_device).manual_seed(seed) |
|
|
output = self( |
|
|
prompt=prompt, |
|
|
height=height, |
|
|
width=width, |
|
|
num_inference_steps=num_sampling_steps, |
|
|
guidance_scale=guidance_scale, |
|
|
num_images_per_prompt=num_images, |
|
|
generator=generator, |
|
|
output_type="pil", |
|
|
return_dict=True, |
|
|
show_progress_bar=True, |
|
|
) |
|
|
return output.images |
|
|
|
