Update pipeline.py

pipeline.py

@@ -660,14 +660,14 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
 
     def __init__(
         self,
-        image_size=512,  # Force 512 resolution
-        in_channels=4,
-        model_channels=320,
-        out_channels=4,
-        num_res_blocks=[2, 2, 2, 2],
-        attention_resolutions=[8, 4, 2],  # Adjusted for 512x512
-        channel_mult=[1, 2, 4, 8],
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
         dropout=0,
+        channel_mult=(1, 2, 4, 8),
         conv_resample=True,
         dims=2,
         num_classes=None,
@@ -688,15 +688,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
     ):
         super().__init__()
         assert context_dim is not None
-        
-        # Add resolution validation
-        assert image_size in [256, 512], "Only 256/512 resolutions supported"
-        super().__init__()
-        
-        # Modify attention resolutions for 512
-        if image_size == 512:
-            attention_resolutions = [16, 8, 4]
-            
+
         if num_heads_upsample == -1:
             num_heads_upsample = num_heads
 
@@ -738,7 +730,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                 f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
                 f"attention will still not be set."
             )
-        
+
         self.attention_resolutions = attention_resolutions
         self.dropout = dropout
         self.channel_mult = channel_mult
@@ -1426,18 +1418,20 @@ class MVDreamPipeline(DiffusionPipeline):
         return torch.zeros_like(image_embeds), image_embeds
 
     def encode_image_latents(self, image, device, num_images_per_prompt):
+
         dtype = next(self.image_encoder.parameters()).dtype
-    
-        # Change interpolation size to match target resolution
-        image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device)
+
+        image = (
+            torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device)
+        )  # [1, 3, H, W]
         image = 2 * image - 1
-        image = F.interpolate(image, (512, 512), mode='bilinear', align_corners=False)  # Changed from 256
+        image = F.interpolate(image, (256, 256), mode="bilinear", align_corners=False)
         image = image.to(dtype=dtype)
 
         posterior = self.vae.encode(image).latent_dist
-        latents = posterior.sample() * self.vae.config.scaling_factor
+        latents = posterior.sample() * self.vae.config.scaling_factor  # [B, C, H, W]
         latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
-    
+
         return torch.zeros_like(latents), latents
 
     @torch.no_grad()
@@ -1445,8 +1439,8 @@ class MVDreamPipeline(DiffusionPipeline):
         self,
         prompt: str = "",
         image: Optional[np.ndarray] = None,
-        height: int = 512,
-        width: int = 512,
+        height: int = 256,
+        width: int = 256,
         elevation: float = 0,
         num_inference_steps: int = 50,
         guidance_scale: float = 7.0,
@@ -1460,10 +1454,6 @@ class MVDreamPipeline(DiffusionPipeline):
         num_frames: int = 4,
         device=torch.device("cuda:0"),
     ):
-        # Add resolution validation
-        if height != 512 or width != 512:
-            raise ValueError("Current implementation requires 512x512 resolution")
-            
         self.unet = self.unet.to(device=device)
         self.vae = self.vae.to(device=device)
         self.text_encoder = self.text_encoder.to(device=device)