rescale noise cfg added

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -48,6 +48,19 @@ class LTX2PipelineOutput:
     frames: jax.Array
     audio: Optional[jax.Array] = None
 
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure.
+    Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891).
+    """
+    std_text = jnp.std(noise_pred_text, axis=list(range(1, noise_pred_text.ndim)), keepdims=True)
+    std_cfg = jnp.std(noise_cfg, axis=list(range(1, noise_cfg.ndim)), keepdims=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
 logger = logging.get_logger(__name__)
 
 
@@ -741,18 +754,25 @@ def __call__(
       num_frames: int = 121,
       frame_rate: float = 24.0,
       num_inference_steps: int = 40,
+      sigmas: Optional[List[float]] = None,
+      timesteps: List[int] = None,
       guidance_scale: float = 3.0,
+      guidance_rescale: float = 0.0,
       noise_scale: float = 1.0,
       num_videos_per_prompt: Optional[int] = 1,
       generator: Optional[jax.Array] = None,
       latents: Optional[jax.Array] = None,
+      audio_latents: Optional[jax.Array] = None,
       prompt_embeds: Optional[jax.Array] = None,
       negative_prompt_embeds: Optional[jax.Array] = None,
       prompt_attention_mask: Optional[jax.Array] = None,
       negative_prompt_attention_mask: Optional[jax.Array] = None,
+      decode_timestep: Union[float, List[float]] = 0.0,
+      decode_noise_scale: Optional[Union[float, List[float]]] = None,
       max_sequence_length: int = 1024,
       dtype: Optional[jnp.dtype] = jnp.float32,
       output_type: str = "pil",
+      return_dict: bool = True,
   ):
       # 1. Check inputs
       self.check_inputs(prompt, height, width, prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask)
@@ -842,7 +862,7 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
            connectors_graphdef, connectors_state, prompt_embeds_jax, prompt_attention_mask_jax.astype(jnp.bool_)
       )
 
-      for i, t in enumerate(tqdm(timesteps)):
+      for i, t in enumerate(timesteps):
           noise_pred, noise_pred_audio = transformer_forward_pass(
               graphdef, state,
               latents_jax,