set num_views as attr of attn_processor to support torch.compile

pipeline_imagedream.py

@@ -76,7 +76,7 @@ class ImageDreamPipeline(StableDiffusionPipeline):
         weight_name: Union[str, List[str]] = "ip-adapter-plus_imagedream.bin",
         image_encoder_folder: Optional[str] = "image_encoder",
         **kwargs,
-    ):
+    ) -> None:
         super().load_ip_adapter(
             pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
             subfolder=subfolder,
@@ -89,12 +89,17 @@ class ImageDreamPipeline(StableDiffusionPipeline):
         if weight_name == "ip-adapter-plus_imagedream.bin":
             setattr(self.image_encoder, "visual_projection", nn.Identity())
             add_imagedream_attn_processor(self.unet)
+            set_num_views(self.unet, self.num_views + 1)
             logging.set_verbosity_error()
             print(
                 "ImageDream Cross-Attention uses `num_views` kwarg, "
                 "and set logging verbosity to error."
             )
 
+    def unload_ip_adapter(self) -> None:
+        super().unload_ip_adapter()
+        set_num_views(self.unet, self.num_views)
+
     def encode_image_to_latents(
         self,
         image: PipelineImageInput,
@@ -326,9 +331,7 @@ class ImageDreamPipeline(StableDiffusionPipeline):
                 guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
             ).to(device=device, dtype=latents.dtype)
 
-        cross_attention_kwargs = {"num_views": num_views}
-        if self.cross_attention_kwargs is not None:
-            cross_attention_kwargs.update(self.cross_attention_kwargs)
+        set_num_views(self.unet, num_views)
 
         # fmt: off
         # 7. Denoising loop
@@ -352,7 +355,7 @@ class ImageDreamPipeline(StableDiffusionPipeline):
                     class_labels=camera,
                     encoder_hidden_states=prompt_embeds,
                     timestep_cond=timestep_cond,
-                    cross_attention_kwargs=cross_attention_kwargs,
+                    cross_attention_kwargs=self.cross_attention_kwargs,
                     added_cond_kwargs=added_cond_kwargs,
                     return_dict=False,
                 )[0]
@@ -508,7 +511,7 @@ def get_camera(
 # fmt: on
 
 
-def add_imagedream_attn_processor(unet: UNet2DConditionModel) -> nn.Module:
+def add_imagedream_attn_processor(unet: UNet2DConditionModel) -> UNet2DConditionModel:
     attn_procs = {}
     for key, attn_processor in unet.attn_processors.items():
         if "attn1" in key:
@@ -519,7 +522,18 @@ def add_imagedream_attn_processor(unet: UNet2DConditionModel) -> nn.Module:
     return unet
 
 
+def set_num_views(unet: UNet2DConditionModel, num_views: int) -> UNet2DConditionModel:
+    for key, attn_processor in unet.attn_processors.items():
+        if isinstance(attn_processor, ImageDreamAttnProcessor2_0):
+            attn_processor.num_views = num_views
+    return unet
+
+
 class ImageDreamAttnProcessor2_0(AttnProcessor2_0):
+    def __init__(self, num_views: int = 4):
+        super().__init__()
+        self.num_views = num_views
+
     def __call__(
         self,
         attn: Attention,
@@ -527,11 +541,10 @@ class ImageDreamAttnProcessor2_0(AttnProcessor2_0):
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         temb: Optional[torch.Tensor] = None,
-        num_views: int = 1,
         *args,
         **kwargs,
     ):
-        if num_views == 1:
+        if self.num_views == 1:
             return super().__call__(
                 attn=attn,
                 hidden_states=hidden_states,
@@ -544,11 +557,11 @@ class ImageDreamAttnProcessor2_0(AttnProcessor2_0):
 
         input_ndim = hidden_states.ndim
         B = hidden_states.size(0)
-        if B % num_views:
+        if B % self.num_views:
             raise ValueError(
-                f"`batch_size`(got {B}) must be a multiple of `num_views`(got {num_views})."
+                f"`batch_size`(got {B}) must be a multiple of `num_views`(got {self.num_views})."
             )
-        real_B = B // num_views
+        real_B = B // self.num_views
         if input_ndim == 4:
             H, W = hidden_states.shape[2:]
             hidden_states = hidden_states.reshape(real_B, -1, H, W).transpose(1, 2)