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