:bug: fix direct load from pipeline, thanks @reynoldscem

scripts/attention.py

@@ -12,8 +12,8 @@ from typing import Optional, Any
 from util import checkpoint
 
 try:
-    import xformers
-    import xformers.ops
+    import xformers # type: ignore
+    import xformers.ops # type: ignore
     XFORMERS_IS_AVAILBLE = True
 except:
     XFORMERS_IS_AVAILBLE = False

scripts/convert_mvdream_to_diffusers.py

@@ -9,13 +9,13 @@ sys.path.insert(0, '../')
 from diffusers.models import (
     AutoencoderKL,
 )
+from omegaconf import OmegaConf
 from diffusers.schedulers import DDIMScheduler
 from diffusers.utils import logging
-
+from typing import Any
 from accelerate import init_empty_weights
 from accelerate.utils import set_module_tensor_to_device
-from rich import print, print_json
-from models import MultiViewUNetModel, MultiViewUNetWrapperModel
+from models import MultiViewUNetWrapperModel
 from pipeline_mvdream import MVDreamStableDiffusionPipeline
 from transformers import CLIPTokenizer, CLIPTextModel
 
@@ -259,14 +259,14 @@ def conv_attn_to_linear(checkpoint):
             if checkpoint[key].ndim > 2:
                 checkpoint[key] = checkpoint[key][:, :, 0]
 
+def create_unet_config(original_config) -> Any:
+    return OmegaConf.to_container(original_config.model.params.unet_config.params, resolve=True)
 
 def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, device):
     checkpoint = torch.load(checkpoint_path, map_location=device)
     # print(f"Checkpoint: {checkpoint.keys()}")
     torch.cuda.empty_cache()
 
-    from omegaconf import OmegaConf
-
     original_config = OmegaConf.load(original_config_file)
     # print(f"Original Config: {original_config}")
     prediction_type = "epsilon"
@@ -296,11 +296,13 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
     #     checkpoint, unet_config, path=None, extract_ema=extract_ema
     # )
     # print(f"Unet Config: {original_config.model.params.unet_config.params}")
-    unet: MultiViewUNetWrapperModel = MultiViewUNetWrapperModel(**original_config.model.params.unet_config.params)
+    unet_config = create_unet_config(original_config)
+    unet: MultiViewUNetWrapperModel = MultiViewUNetWrapperModel(**unet_config)
+    unet.register_to_config(**unet_config)
     # print(f"Unet State Dict: {unet.state_dict().keys()}")
     unet.load_state_dict({key.replace("model.diffusion_model.", "unet."): value for key, value in checkpoint.items() if key.replace("model.diffusion_model.", "unet.") in unet.state_dict()})
     for param_name, param in unet.state_dict().items():
-        set_module_tensor_to_device(unet, param_name, "cuda:0", value=param)
+        set_module_tensor_to_device(unet, param_name, device=device, value=param)
 
     # Convert the VAE model.
     vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
@@ -316,18 +318,18 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
     with init_empty_weights():
         vae = AutoencoderKL(**vae_config)
 
+    for param_name, param in converted_vae_checkpoint.items():
+        set_module_tensor_to_device(vae, param_name, device=device, value=param)
+
     if original_config.model.params.unet_config.params.context_dim == 768:
         tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device=torch.device("cuda:0")) # type: ignore
+        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device=device) # type: ignore
     elif original_config.model.params.unet_config.params.context_dim == 1024:
         tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="tokenizer")
-        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=torch.device("cuda:0")) # type: ignore
+        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=device) # type: ignore
     else:
         raise ValueError(f"Unknown context_dim: {original_config.model.paams.unet_config.params.context_dim}")
 
-    for param_name, param in converted_vae_checkpoint.items():
-        set_module_tensor_to_device(vae, param_name, "cuda:0", value=param)
-
     pipe = MVDreamStableDiffusionPipeline(
         vae=vae,
         unet=unet,
@@ -359,6 +361,8 @@ if __name__ == "__main__":
     parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
     parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)")
     args = parser.parse_args()
+    
+    args.device = torch.device(args.device if args.device is not None else "cuda" if torch.cuda.is_available() else "cpu")
 
     pipe = convert_from_original_mvdream_ckpt(
         checkpoint_path=args.checkpoint_path,
@@ -369,15 +373,36 @@ if __name__ == "__main__":
     if args.half:
         pipe.to(torch_dtype=torch.float16)
 
-    if args.test:
-        images = pipe(
-            prompt="Head of Hatsune Miku",
-            negative_prompt="painting, bad quality, flat",
-            output_type="pil",
-            guidance_scale=7.5,
-            num_inference_steps=50,
-        )
-        for i, image in enumerate(images):
-            image.save(f"image_{i}.png") # type: ignore
-
+    print(f"Saving pipeline to {args.dump_path}...")
     pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
+    
+    if args.test:
+        try:
+            print(f"Testing each subcomponent of the pipeline...")
+            images = pipe(
+                prompt="Head of Hatsune Miku",
+                negative_prompt="painting, bad quality, flat",
+                output_type="pil",
+                guidance_scale=7.5,
+                num_inference_steps=50,
+                device=args.device,
+            )
+            for i, image in enumerate(images):
+                image.save(f"image_{i}.png") # type: ignore
+
+            print(f"Testing entire pipeline...")
+            loaded_pipe: MVDreamStableDiffusionPipeline = MVDreamStableDiffusionPipeline.from_pretrained(args.dump_path, safe_serialization=args.to_safetensors) # type: ignore
+            images = loaded_pipe(
+                prompt="Head of Hatsune Miku",
+                negative_prompt="painting, bad quality, flat",
+                output_type="pil",
+                guidance_scale=7.5,
+                num_inference_steps=50,
+                device=args.device,
+            )
+            for i, image in enumerate(images):
+                image.save(f"image_{i}.png") # type: ignore
+        except Exception as e:
+            print(f"Failed to test inference: {e}")
+            raise e from e
+        print("Inference test passed!")

scripts/models.py

@@ -25,9 +25,72 @@ from torch import Tensor
 
 class MultiViewUNetWrapperModel(ModelMixin, ConfigMixin):
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, 
+            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,
+            use_checkpoint=False,
+            use_fp16=False,
+            use_bf16=False,
+            num_heads=-1,
+            num_head_channels=-1,
+            num_heads_upsample=-1,
+            use_scale_shift_norm=False,
+            resblock_updown=False,
+            use_new_attention_order=False,
+            use_spatial_transformer=False, # custom transformer support
+            transformer_depth=1, # custom transformer support
+            context_dim=None, # custom transformer support
+            n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
+            legacy=True,
+            disable_self_attentions=None,
+            num_attention_blocks=None,
+            disable_middle_self_attn=False,
+            use_linear_in_transformer=False,
+            adm_in_channels=None,
+            camera_dim=None,):
         super().__init__()
-        self.unet: MultiViewUNetModel = MultiViewUNetModel(*args, **kwargs)
+        self.unet: MultiViewUNetModel = MultiViewUNetModel(
+            image_size=image_size,
+            in_channels=in_channels,
+            model_channels=model_channels,
+            out_channels=out_channels,
+            num_res_blocks=num_res_blocks,
+            attention_resolutions=attention_resolutions,
+            dropout=dropout,
+            channel_mult=channel_mult,
+            conv_resample=conv_resample,
+            dims=dims,
+            num_classes=num_classes,
+            use_checkpoint=use_checkpoint,
+            use_fp16=use_fp16,
+            use_bf16=use_bf16,
+            num_heads=num_heads,
+            num_head_channels=num_head_channels,
+            num_heads_upsample=num_heads_upsample,
+            use_scale_shift_norm=use_scale_shift_norm,
+            resblock_updown=resblock_updown,
+            use_new_attention_order=use_new_attention_order,
+            use_spatial_transformer=use_spatial_transformer,
+            transformer_depth=transformer_depth,
+            context_dim=context_dim,
+            n_embed=n_embed,
+            legacy=legacy,
+            disable_self_attentions=disable_self_attentions,
+            num_attention_blocks=num_attention_blocks,
+            disable_middle_self_attn=disable_middle_self_attn,
+            use_linear_in_transformer=use_linear_in_transformer,
+            adm_in_channels=adm_in_channels,
+            camera_dim=camera_dim,
+        )
 
     def forward(self, *args, **kwargs):
         return self.unet(*args, **kwargs)

scripts/pipeline_mvdream.py

@@ -1,16 +1,14 @@
 import torch
 import numpy as np
 import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Callable, List, Optional, Union
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers import AutoencoderKL, DiffusionPipeline
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.utils import (
     deprecate,
     is_accelerate_available,
     is_accelerate_version,
     logging,
-    replace_example_docstring,
 )
 from diffusers.configuration_utils import FrozenDict
 from diffusers.schedulers import DDIMScheduler
@@ -20,6 +18,7 @@ except ImportError:
     from diffusers.utils.torch_utils import randn_tensor # new import # type: ignore
 
 from models import MultiViewUNetWrapperModel
+from accelerate.utils import set_module_tensor_to_device
 
 logger = logging.get_logger(__name__) # pylint: disable=invalid-name
 
@@ -391,9 +390,13 @@ class MVDreamStableDiffusionPipeline(DiffusionPipeline):
         output_type: Optional[str] = "pil",
         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
         callback_steps: int = 1,
+        batch_size: int = 4,
+        device = torch.device("cuda:0"),
     ):
-        batch_size = 4
-        device = torch.device("cuda:0")
+        self.unet = self.unet.to(device=device)
+        self.vae = self.vae.to(device=device)
+
+        self.text_encoder = self.text_encoder.to(device=device)
 
         camera = get_camera(batch_size).to(device=device)
 

unet/config.json

@@ -1,4 +1,27 @@
 {
   "_class_name": "MultiViewUNetWrapperModel",
-  "_diffusers_version": "0.21.4"
+  "_diffusers_version": "0.21.4",
+  "attention_resolutions": [
+    4,
+    2,
+    1
+  ],
+  "camera_dim": 16,
+  "channel_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "context_dim": 768,
+  "image_size": 32,
+  "in_channels": 4,
+  "legacy": false,
+  "model_channels": 320,
+  "num_heads": 8,
+  "num_res_blocks": 2,
+  "out_channels": 4,
+  "transformer_depth": 1,
+  "use_checkpoint": false,
+  "use_spatial_transformer": true
 }