i

apps/mv_models.py

@@ -20,98 +20,105 @@ import spaces
 
 parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-
+    
 class GenMVImage(object):
     def __init__(self, device):
         self.seed = 1024
         self.guidance_scale = 7.5
         self.step = 50
+        self.pipelines = {}
         self.device = device
-
+    
+    @spaces.GPU
+    def gen_image_from_crm(self, image):
         from .third_party.CRM.pipelines import TwoStagePipeline
         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
         stage1_model_config.resume = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth", repo_type="model")
         stage1_model_config.config = f"{parent_dir}/apps/third_party/CRM/" + stage1_model_config.config
-        self.crm_pipeline = TwoStagePipeline(
-                                    stage1_model_config,
-                                    stage1_sampler_config,
-                                    device=self.device,
-                                    dtype=torch.float16
-                                )
-        self.crm_pipeline.set_seed(self.seed)
-
-        # 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
-        # self.wonder3d_pipeline =  DiffusionPipeline.from_pretrained(
-        #     'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
-        #     custom_pipeline='flamehaze1115/wonder3d-pipeline',
-        #     torch_dtype=torch.float16
-        # )
-        # self.wonder3d_pipeline.unet.enable_xformers_memory_efficient_attention()
-        # self.wonder3d_pipeline.to(self.device)
-        # self.wonder3d_pipeline.set_progress_bar_config(disable=True)
-
-
-        # sys.path.append(f"{parent_dir}/apps/third_party/mvdream_diffusers")
-        # from .third_party.mvdream_diffusers.pipeline_mvdream import MVDreamPipeline
-        # self.mvdream_pipeline = MVDreamPipeline.from_pretrained(
-        #     "ashawkey/mvdream-sd2.1-diffusers", # remote weights
-        #     torch_dtype=torch.float16,
-        #     trust_remote_code=True,
-        # )
-        # self.mvdream_pipeline = self.mvdream_pipeline.to(self.device)
-        # self.imagedream_pipeline = MVDreamPipeline.from_pretrained(
-        #         "ashawkey/imagedream-ipmv-diffusers", # remote weights
-        #         torch_dtype=torch.float16,
-        #         trust_remote_code=True,
-        #     )
-        # self.imagedream_pipeline = self.imagedream_pipeline.to(self.device)
-
-
-    @spaces.GPU
-    def gen_image_from_crm(self, image):
-        rt_dict = self.crm_pipeline(
-            image, 
-            scale=self.guidance_scale, 
-            step=self.step
-        )
+        if "crm" in self.pipelines.keys():
+            pipeline = self.pipelines['crm']
+        else:
+            self.pipelines['crm'] = TwoStagePipeline(
+                                        stage1_model_config,
+                                        stage1_sampler_config,
+                                        device=self.device,
+                                        dtype=torch.float16
+                                    )
+            pipeline = self.pipelines['crm']
+        pipeline.set_seed(self.seed)
+        rt_dict = pipeline(image, scale=self.guidance_scale, step=self.step)
         mv_imgs = rt_dict["stage1_images"]
         return mv_imgs[5], mv_imgs[3], mv_imgs[2], mv_imgs[0]
     
     @spaces.GPU
     def gen_image_from_mvdream(self, image, text):
+        from .third_party.mvdream_diffusers.pipeline_mvdream import MVDreamPipeline
         if image is None:
-            mv_imgs = self.mvdream_pipeline(
-                text,
-                negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy",
-                num_inference_steps=self.step,
-                guidance_scale=self.guidance_scale,
-                generator = torch.Generator(self.device).manual_seed(self.seed)
-            )
-        elif text is not None:
+            if "mvdream" in self.pipelines.keys():
+                pipe_MVDream = self.pipelines['mvdream']
+            else:
+                self.pipelines['mvdream'] = MVDreamPipeline.from_pretrained(
+                    "ashawkey/mvdream-sd2.1-diffusers", # remote weights
+                    torch_dtype=torch.float16,
+                    trust_remote_code=True,
+                )
+                self.pipelines['mvdream'] = self.pipelines['mvdream'].to(self.device)
+                pipe_MVDream = self.pipelines['mvdream']
+            mv_imgs = pipe_MVDream(
+                    text,
+                    negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy",
+                    num_inference_steps=self.step,
+                    guidance_scale=self.guidance_scale,
+                    generator = torch.Generator(self.device).manual_seed(self.seed)
+                )
+        else:
             image = np.array(image)
             image = image.astype(np.float32) / 255.0
             image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4])
-
-            mv_imgs = self.imagedream_pipeline(
-                text,
-                image,
-                negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy",
-                num_inference_steps=self.step,
-                guidance_scale=self.guidance_scale,
-                generator = torch.Generator(self.device).manual_seed(self.seed)
-            )
+            if "imagedream" in self.pipelines.keys():
+                pipe_imagedream = self.pipelines['imagedream']
+            else:
+                self.pipelines['imagedream'] = MVDreamPipeline.from_pretrained(
+                        "ashawkey/imagedream-ipmv-diffusers", # remote weights
+                        torch_dtype=torch.float16,
+                        trust_remote_code=True,
+                    )
+                self.pipelines['imagedream'] = self.pipelines['imagedream'].to(self.device)
+                pipe_imagedream = self.pipelines['imagedream']
+            mv_imgs = pipe_imagedream(
+                        text,
+                        image,
+                        negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy",
+                        num_inference_steps=self.step,
+                        guidance_scale=self.guidance_scale,
+                        generator = torch.Generator(self.device).manual_seed(self.seed)
+                    )
         return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
     
     @spaces.GPU
     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
         batch = prepare_data(image, crop_size)
 
-        generator = torch.Generator(device=self.wonder3d_pipeline.unet.device).manual_seed(self.seed)
+        if "wonder3d" in self.pipelines.keys():
+            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
+            )
+            self.pipelines['wonder3d'].unet.enable_xformers_memory_efficient_attention()
+            self.pipelines['wonder3d'].to(self.device)
+            self.pipelines['wonder3d'].set_progress_bar_config(disable=True)
+            pipeline = self.pipelines['wonder3d']
+            
+        generator = torch.Generator(device=pipeline.unet.device).manual_seed(self.seed)
         # repeat  (2B, Nv, 3, H, W)
         imgs_in = torch.cat([batch['imgs_in']] * 2, dim=0).to(weight_dtype)
 
@@ -126,9 +133,8 @@ class GenMVImage(object):
         imgs_in = rearrange(imgs_in, "Nv C H W -> (Nv) C H W")
         # (B*Nv, Nce)
 
-        out = self.wonder3d_pipeline(
+        out = pipeline(
             imgs_in,
-            # camera_embeddings,
             generator=generator,
             guidance_scale=self.guidance_scale,
             num_inference_steps=self.step,
@@ -146,7 +152,7 @@ class GenMVImage(object):
 
         mv_imgs = images_pred
         return mv_imgs[0], mv_imgs[2], mv_imgs[4], mv_imgs[5]
-        
+
     @spaces.GPU
     def run(self, mvimg_model, text, image, crop_size, seed, guidance_scale, step):
         self.seed = seed
@@ -155,6 +161,6 @@ class GenMVImage(object):
         if mvimg_model.upper() == "CRM":
             return self.gen_image_from_crm(image)
         elif mvimg_model.upper() == "IMAGEDREAM":
-            return self.gen_image_from_mvdream(image, None)
+            return self.gen_image_from_mvdream(image, text)
         elif mvimg_model.upper() == "WONDER3D":
             return self.gen_image_from_wonder3d(image, crop_size)

gradio_app.py

@@ -141,8 +141,8 @@ if __name__=="__main__":
         "Auto Remove Background": "Auto Remove Background",
         "Original Image": "Original Image",        
     })
-    # mvimg_model_config_list = ["CRM", "ImageDream", "Wonder3D"]
     mvimg_model_config_list = ["CRM"]
+    # mvimg_model_config_list = ["CRM", "ImageDream", "Wonder3D"]
     
     # for 3D latent set diffusion
     ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt", repo_type="model")