Add initial support for chroma radiance

extensions_built_in/diffusion_models/chroma/pipeline.py

@@ -6,6 +6,7 @@ from diffusers import FluxPipeline
 from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
 from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
 from diffusers.utils import is_torch_xla_available
+from diffusers.utils.torch_utils import randn_tensor
 
 
 if is_torch_xla_available():
@@ -16,7 +17,134 @@ else:
     XLA_AVAILABLE = False
 
 
+def prepare_latent_image_ids(batch_size, height, width, patch_size=2, max_offset=0):
+    """
+    Generates positional embeddings for a latent image.
+
+    Args:
+        batch_size (int): The number of images in the batch.
+        height (int): The height of the image.
+        width (int): The width of the image.
+        patch_size (int, optional): The size of the patches. Defaults to 2.
+        max_offset (int, optional): The maximum random offset to apply. Defaults to 0.
+
+    Returns:
+        torch.Tensor: A tensor containing the positional embeddings.
+    """
+    # the random pos embedding helps generalize to larger res without training at large res
+    # pos embedding for rope, 2d pos embedding, corner embedding and not center based
+    latent_image_ids = torch.zeros(height // patch_size, width // patch_size, 3)
+
+    # Add positional encodings
+    latent_image_ids[..., 1] = (
+        latent_image_ids[..., 1] + torch.arange(height // patch_size)[:, None]
+    )
+    latent_image_ids[..., 2] = (
+        latent_image_ids[..., 2] + torch.arange(width // patch_size)[None, :]
+    )
+
+    # Add random offset if specified
+    if max_offset > 0:
+        offset_y = torch.randint(0, max_offset + 1, (1,)).item()
+        offset_x = torch.randint(0, max_offset + 1, (1,)).item()
+        latent_image_ids[..., 1] += offset_y
+        latent_image_ids[..., 2] += offset_x
+
+
+    (
+        latent_image_id_height,
+        latent_image_id_width,
+        latent_image_id_channels,
+    ) = latent_image_ids.shape
+
+    # Reshape for batch
+    latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
+    latent_image_ids = latent_image_ids.reshape(
+        batch_size,
+        latent_image_id_height * latent_image_id_width,
+        latent_image_id_channels,
+    )
+
+    return latent_image_ids
+
+
 class ChromaPipeline(FluxPipeline):
+    def __init__(
+        self,
+        scheduler,
+        vae,
+        text_encoder,
+        tokenizer,
+        text_encoder_2,
+        tokenizer_2,
+        transformer,
+        image_encoder = None,
+        feature_extractor = None,
+        is_radiance: bool = False,
+    ):
+        super().__init__(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+        )
+        self.is_radiance = is_radiance
+        self.vae_scale_factor = 8 if not is_radiance else 1
+    
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = prepare_latent_image_ids(
+                batch_size, 
+                height, 
+                width, 
+                patch_size=2 if not self.is_radiance else 16
+            ).to(device=device, dtype=dtype)
+            # latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        
+        if not self.is_radiance:
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+        # latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+        latent_image_ids = prepare_latent_image_ids(
+            batch_size, 
+            height, 
+            width, 
+            patch_size=2 if not self.is_radiance else 16
+        ).to(device=device, dtype=dtype)
+
+        return latents, latent_image_ids
+    
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -70,6 +198,8 @@ class ChromaPipeline(FluxPipeline):
 
         # 4. Prepare latent variables
         num_channels_latents = 64 // 4
+        if self.is_radiance:
+            num_channels_latents = 3
         latents, latent_image_ids = self.prepare_latents(
             batch_size * num_images_per_prompt,
             num_channels_latents,
@@ -82,8 +212,8 @@ class ChromaPipeline(FluxPipeline):
         )
         
         # extend img ids to match batch size
-        latent_image_ids = latent_image_ids.unsqueeze(0)
-        latent_image_ids = torch.cat([latent_image_ids] * batch_size, dim=0)
+        # latent_image_ids = latent_image_ids.unsqueeze(0)
+        # latent_image_ids = torch.cat([latent_image_ids] * batch_size, dim=0)
 
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
@@ -180,8 +310,9 @@ class ChromaPipeline(FluxPipeline):
             image = latents
 
         else:
-            latents = self._unpack_latents(
-                latents, height, width, self.vae_scale_factor)
+            if not self.is_radiance:
+                latents = self._unpack_latents(
+                    latents, height, width, self.vae_scale_factor)
             latents = (latents / self.vae.config.scaling_factor) + \
                 self.vae.config.shift_factor
             image = self.vae.decode(latents, return_dict=False)[0]