implements skip layer guidance for better generations.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -207,15 +207,20 @@ prompt: "A cat and a dog baking a cake together in a kitchen. The cat is careful
 prompt_2: "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."
 negative_prompt: "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
 do_classifier_free_guidance: True
-height: 720
-width: 1280
+height: 480
+width: 832
 num_frames: 81
 guidance_scale: 5.0
+flow_shift: 3.0
+
+# skip layer guidance
+slg_layers: [9]
+slg_start: 0.2
+slg_end: 1.0
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0
 num_inference_steps: 30
 save_final_checkpoint: False
-flow_shift: 5.0
 
 # SDXL Lightning parameters
 lightning_from_pt: True

src/maxdiffusion/generate_wan.py

@@ -21,8 +21,15 @@
 from maxdiffusion.utils import export_to_video
 
 def run(config):
+  print("seed: ", config.seed)
   pipeline = WanPipeline.from_pretrained(config)
   s0 = time.perf_counter()
+  
+  # Skip layer guidance
+  slg_layers = config.slg_layers
+  slg_start = config.slg_start
+  slg_end = config.slg_end
+  
   videos = pipeline(
     prompt=config.prompt,
     negative_prompt=config.negative_prompt,
@@ -31,6 +38,9 @@ def run(config):
     num_frames=config.num_frames,
     num_inference_steps=config.num_inference_steps,
     guidance_scale=config.guidance_scale,
+    slg_layers=slg_layers,
+    slg_start=slg_start,
+    slg_end=slg_end
   )
 
   print("compile time: ", (time.perf_counter() - s0))
@@ -46,6 +56,9 @@ def run(config):
       num_frames=config.num_frames,
       num_inference_steps=config.num_inference_steps,
       guidance_scale=config.guidance_scale,
+      slg_layers=slg_layers,
+      slg_start=slg_start,
+      slg_end=slg_end
     )
   print("generation time: ", (time.perf_counter() - s0))
   for i in range(len(videos)):

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -14,7 +14,7 @@
  limitations under the License.
 """
 
-from typing import Tuple, Optional, Dict, Union, Any
+from typing import Tuple, Optional, Dict, Union, Any, List
 import math
 import jax
 import jax.numpy as jnp
@@ -453,6 +453,8 @@ def __call__(
     hidden_states: jax.Array,
     timestep: jax.Array,
     encoder_hidden_states: jax.Array,
+    is_uncond: jax.Array, # jnp.bool_ scalar
+    slg_mask: jax.Array, # jnp.bool_ array of shape (num_blocks,)
     encoder_hidden_states_image: Optional[jax.Array] = None,
     return_dict: bool = True,
     attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -476,8 +478,14 @@ def __call__(
 
     if encoder_hidden_states_image is not None:
       raise NotImplementedError("img2vid is not yet implemented.")
-    for block in self.blocks:
-      hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
+    for block_idx, block in enumerate(self.blocks):
+      should_skip_block = slg_mask[block_idx] & is_uncond
+      hidden_states = jax.lax.cond(
+        should_skip_block,
+        lambda hs: hs, # If true, pass through original hidden_states (skip block)
+        lambda _: block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb),
+        hidden_states
+      )
     shift, scale = jnp.split(self.scale_shift_table + jnp.expand_dims(temb, axis=1), 2, axis=1)
 
     hidden_states = (self.norm_out(hidden_states.astype(jnp.float32)) * (1 + scale) + shift).astype(hidden_states.dtype)

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -369,7 +369,10 @@ def __call__(
     latents: jax.Array = None,
     prompt_embeds: jax.Array = None,
     negative_prompt_embeds: jax.Array = None,
-    vae_only: bool = False
+    vae_only: bool = False,
+    slg_layers: List[int] = None,
+    slg_start: float = 0.0,
+    slg_end: float = 1.0
   ):
     if not vae_only:
       if num_frames % self.vae_scale_factor_temporal != 1:
@@ -424,7 +427,11 @@ def __call__(
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         scheduler=self.scheduler,
-        scheduler_state=scheduler_state
+        scheduler_state=scheduler_state,
+        slg_layers=slg_layers,
+        slg_start=slg_start,
+        slg_end=slg_end,
+        num_transformer_layers=self.transformer.config.num_layers
       )
 
       with self.mesh:
@@ -450,12 +457,22 @@ def __call__(
 
 
 @jax.jit
-def transformer_forward_pass(graphdef, sharded_state, rest_of_state, latents, timestep, prompt_embeds):
+def transformer_forward_pass(
+  graphdef,
+  sharded_state,
+  rest_of_state,
+  latents,
+  timestep,
+  prompt_embeds,
+  is_uncond,
+  slg_mask):
   wan_transformer = nnx.merge(graphdef, sharded_state, rest_of_state)
   return wan_transformer(
     hidden_states=latents,
     timestep=timestep,
-    encoder_hidden_states=prompt_embeds
+    encoder_hidden_states=prompt_embeds,
+    is_uncond=is_uncond,
+    slg_mask=slg_mask
   )[0]
 
 #@partial(jax.jit, static_argnums=(6, 7, 8))
@@ -469,16 +486,42 @@ def run_inference(
   guidance_scale: float,
   num_inference_steps: int,
   scheduler : FlaxUniPCMultistepScheduler,
-  scheduler_state):
+  num_transformer_layers: int,
+  scheduler_state,
+  slg_layers: List[int] = None,
+  slg_start: float = 0.0,
+  slg_end: float = 1.0
+  ):
     do_classifier_free_guidance = guidance_scale > 1.0
     for step in range(num_inference_steps):
+      slg_mask = jnp.zeros(num_transformer_layers, dtype=jnp.bool_)
+      if slg_layers and int(slg_start * num_inference_steps) <= step < int(slg_end * num_inference_steps):
+        slg_mask = slg_mask.at[jnp.array(slg_layers)].set(True)
       t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
       timestep = jnp.broadcast_to(t, latents.shape[0])
 
-      noise_pred = transformer_forward_pass(graphdef, sharded_state, rest_of_state, latents, timestep, prompt_embeds)
+      noise_pred = transformer_forward_pass(
+        graphdef,
+        sharded_state,
+        rest_of_state,
+        latents,
+        timestep,
+        prompt_embeds,
+        is_uncond=jnp.array(False, dtype=jnp.bool_),
+        slg_mask=slg_mask
+      )
 
       if do_classifier_free_guidance:
-        noise_uncond = transformer_forward_pass(graphdef, sharded_state, rest_of_state, latents, timestep, negative_prompt_embeds)
+        noise_uncond = transformer_forward_pass(
+          graphdef,
+          sharded_state,
+          rest_of_state,
+          latents,
+          timestep,
+          negative_prompt_embeds,
+          is_uncond=jnp.array(True, dtype=jnp.bool_),
+          slg_mask=slg_mask
+          )
         noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
     return latents