add missing imports

README.md

@@ -1,4 +1,4 @@
-# MVDream - 3D Generation
+# MVDream - threestudio
 Yichun Shi, Peng Wang, Jianglong Ye, Long Mai, Kejie Li, Xiao Yang
 
 | [Project Page](https://mv-dream.github.io/) | [Paper](https://arxiv.org/abs/2308.16512) |

requirements.txt

@@ -39,3 +39,6 @@ git+https://github.com/openai/CLIP.git
 
 #controlnet
 controlnet_aux
+
+# mvdream
+open-clip-torch==2.7.0

threestudio/data/__init__.py

@@ -1 +1 @@
-from . import co3d, image, multiview, uncond
+from . import co3d, image, multiview, uncond, random_multiview

threestudio/models/background/neural_environment_map_background.py

@@ -31,6 +31,7 @@ class Config(BaseBackground.Config):
         )
         random_aug: bool = False
         random_aug_prob: float = 0.5
+        share_aug_bg: bool = False
         eval_color: Optional[Tuple[float, float, float]] = None
 
     cfg: Config
@@ -59,8 +60,9 @@ def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
             and random.random() < self.cfg.random_aug_prob
         ):
             # use random background color with probability random_aug_prob
+            n_color = 1 if self.cfg.share_aug_bg else dirs.shape[0]
             color = color * 0 + (  # prevent checking for unused parameters in DDP
-                torch.rand(dirs.shape[0], 1, 1, self.cfg.n_output_dims)
+                torch.rand(n_color, 1, 1, self.cfg.n_output_dims)
                 .to(dirs)
                 .expand(*dirs.shape[:-1], -1)
             )

threestudio/models/guidance/__init__.py

@@ -5,4 +5,5 @@
     stable_diffusion_guidance,
     stable_diffusion_vsd_guidance,
     zero123_guidance,
+    multiview_diffusion_guidance,
 )

threestudio/models/prompt_processors/stable_diffusion_prompt_processor.py

@@ -88,7 +88,7 @@ def spawn_func(pretrained_model_name_or_path, prompts, cache_dir):
                 max_length=tokenizer.model_max_length,
                 return_tensors="pt",
             )
-            text_embeddings = text_encoder(tokens.input_ids)[0]
+            text_embeddings = text_encoder(tokens.input_ids.to(text_encoder.device))[0]
 
         for prompt, embedding in zip(prompts, text_embeddings):
             torch.save(

threestudio/systems/__init__.py

@@ -10,4 +10,5 @@
     sjc,
     textmesh,
     zero123,
+    mvdream,
 )