i

apps/mv_models.py

@@ -19,36 +19,6 @@ from huggingface_hub import hf_hub_download
 
 parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-    
-@dataclass
-class TestConfig:
-    pretrained_model_name_or_path: str
-    pretrained_unet_path: str
-    revision: Optional[str]
-    validation_dataset: Dict
-    save_dir: str
-    seed: Optional[int]
-    validation_batch_size: int
-    dataloader_num_workers: int
-
-    local_rank: int
-
-    pipe_kwargs: Dict
-    pipe_validation_kwargs: Dict
-    unet_from_pretrained_kwargs: Dict
-    validation_guidance_scales: List[float]
-    validation_grid_nrow: int
-    camera_embedding_lr_mult: float
-
-    num_views: int
-    camera_embedding_type: str
-
-    pred_type: str  # joint, or ablation
-
-    enable_xformers_memory_efficient_attention: bool
-
-    cond_on_normals: bool
-    cond_on_colors: bool
 
 class GenMVImage(object):
     def __init__(self, device):
@@ -59,9 +29,7 @@ class GenMVImage(object):
         self.device = device
         
     def gen_image_from_crm(self, image):
-        
         from .third_party.CRM.pipelines import TwoStagePipeline
-        specs = json.load(open(f"{parent_dir}/apps/third_party/CRM/configs/specs_objaverse_total.json"))
         stage1_config = OmegaConf.load(f"{parent_dir}/apps/third_party/CRM/configs/nf7_v3_SNR_rd_size_stroke.yaml").config
         stage1_sampler_config = stage1_config.sampler
         stage1_model_config = stage1_config.models
@@ -127,6 +95,7 @@ class GenMVImage(object):
         return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
     
     def gen_image_from_wonder3d(self, image, crop_size):
+        sys.path.append(f"{parent_dir}/apps/third_party/Wonder3D")
         from diffusers import DiffusionPipeline  # only tested on diffusers[torch]==0.19.3, may have conflicts with newer versions of diffusers
 
         weight_dtype = torch.float16
@@ -136,9 +105,9 @@ class GenMVImage(object):
             pipeline = self.pipelines['wonder3d']
         else:
             self.pipelines['wonder3d'] =  DiffusionPipeline.from_pretrained(
-            'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
-            custom_pipeline='flamehaze1115/wonder3d-pipeline',
-            torch_dtype=torch.float16
+                'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
+                custom_pipeline='flamehaze1115/wonder3d-pipeline',
+                torch_dtype=torch.float16
             )
             self.pipelines['wonder3d'].unet.enable_xformers_memory_efficient_attention()
             self.pipelines['wonder3d'].to(self.device)