Rename pipeline_imagedream.py to pipeline_mvdiffusion.py

pipeline_imagedream.py → pipeline_mvdiffusion.py

@@ -41,7 +41,7 @@ from transformers import (
 )
 
 
-class ImageDreamPipeline(StableDiffusionPipeline):
+class MVDiffusionPipeline(StableDiffusionPipeline):
     def __init__(
         self,
         vae: AutoencoderKL,
@@ -50,15 +50,15 @@ class ImageDreamPipeline(StableDiffusionPipeline):
         unet: UNet2DConditionModel,
         scheduler: KarrasDiffusionSchedulers,
         safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
-        image_encoder: CLIPVisionModel = None,
+        feature_extractor: Optional[CLIPImageProcessor] = None,
+        image_encoder: Optional[CLIPVisionModel] = None,
         requires_safety_checker: bool = False,
     ) -> None:
         super().__init__(
             vae=vae,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-            unet=add_imagedream_attn_processor(unet),
+            unet=add_mv_attn_processor(unet),
             scheduler=scheduler,
             safety_checker=safety_checker,
             feature_extractor=feature_extractor,
@@ -88,7 +88,7 @@ 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)
+            add_mv_attn_processor(self.unet)
             set_num_views(self.unet, self.num_views + 1)
 
     def unload_ip_adapter(self) -> None:
@@ -193,7 +193,7 @@ class ImageDreamPipeline(StableDiffusionPipeline):
         if cross_attention_kwargs is None:
             num_views = self.num_views
         else:
-            num_views = cross_attention_kwargs.pop("num_views", self.num_views)
+            cross_attention_kwargs.pop("num_views", self.num_views)
 
         # 0. Default height and width to unet
         height = height or self.unet.config.sample_size * self.vae_scale_factor
@@ -506,11 +506,11 @@ def get_camera(
 # fmt: on
 
 
-def add_imagedream_attn_processor(unet: UNet2DConditionModel) -> UNet2DConditionModel:
+def add_mv_attn_processor(unet: UNet2DConditionModel, num_views: int = 4) -> UNet2DConditionModel:
     attn_procs = {}
     for key, attn_processor in unet.attn_processors.items():
         if "attn1" in key:
-            attn_procs[key] = ImageDreamAttnProcessor2_0()
+            attn_procs[key] = MVAttnProcessor2_0(num_views)
         else:
             attn_procs[key] = attn_processor
     unet.set_attn_processor(attn_procs)
@@ -519,12 +519,12 @@ def add_imagedream_attn_processor(unet: UNet2DConditionModel) -> UNet2DCondition
 
 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):
+        if isinstance(attn_processor, MVAttnProcessor2_0):
             attn_processor.num_views = num_views
     return unet
 
 
-class ImageDreamAttnProcessor2_0(AttnProcessor2_0):
+class MVAttnProcessor2_0(AttnProcessor2_0):
     def __init__(self, num_views: int = 4):
         super().__init__()
         self.num_views = num_views