update

Author: ishan-modi

src/diffusers/models/controlnets/controlnet_sana.py

@@ -40,6 +40,8 @@ class SanaControlNetOutput(BaseOutput):
 
 class SanaControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["SanaTransformerBlock", "PatchEmbed"]
+    _skip_layerwise_casting_patterns = ["patch_embed", "norm"]
 
     @register_to_config
     def __init__(

src/diffusers/pipelines/sana/pipeline_sana_controlnet.py

@@ -207,8 +207,8 @@ class SanaControlNetPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
     bad_punct_regex = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}")
     # fmt: on
 
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+    model_cpu_offload_seq = "text_encoder->controlnet->transformer->vae"
+    _callback_tensor_inputs = ["latents", "control_image", "prompt_embeds", "negative_prompt_embeds"]
 
     def __init__(
         self,