Update gradio_app.py

gradio_app.py

@@ -14,6 +14,7 @@ from collections import OrderedDict
 import trimesh
 import gradio as gr
 from typing import Any
+from einops import rearrange
 
 proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.append(os.path.join(proj_dir))
@@ -58,6 +59,8 @@ If you have any questions, feel free to open a discussion or contact us at <b>we
 """
 from apps.third_party.CRM.pipelines import TwoStagePipeline
 from apps.third_party.LGM.pipeline_mvdream import MVDreamPipeline
+from apps.third_party.Era3D.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline
+from apps.third_party.Era3D.data.single_image_dataset import SingleImageDataset 
 
 import re
 import os
@@ -88,22 +91,25 @@ chmod(f"{parent_dir}/apps/third_party/InstantMeshes", "777")
 
 model = None
 cached_dir = None
-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
+generator = None
+
+sys.path.append(f"apps/third_party/CRM")
 crm_pipeline = None
 
 sys.path.append(f"apps/third_party/LGM")
 imgaedream_pipeline = None
 
+sys.path.append(f"apps/third_party/Era3D")
+era3d_pipeline = None
+
 @spaces.GPU
 def gen_mvimg(
     mvimg_model, image, seed, guidance_scale, step, text, neg_text, elevation, backgroud_color
 ):
     if seed == 0:
         seed = np.random.randint(1, 65535)
+        global generator
+        generator.manual_seed(seed)
 
     if mvimg_model == "CRM":
         global crm_pipeline
@@ -118,7 +124,7 @@ def gen_mvimg(
         return mv_imgs[5], mv_imgs[3], mv_imgs[2], mv_imgs[0]
     
     elif mvimg_model == "ImageDream":
-        global imagedream_pipeline, generator
+        global imagedream_pipeline
         background = Image.new("RGBA", image.size, backgroud_color)
         image = Image.alpha_composite(background, image)
         image = np.array(image).astype(np.float32) / 255.0
@@ -130,9 +136,36 @@ def gen_mvimg(
             guidance_scale=guidance_scale,  
             num_inference_steps=step, 
             elevation=elevation,
+            generator=generator,
         )
         return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
-    
+
+    elif mvimg_model == "Era3D":
+        global era3d_pipeline
+        crop_size = 420
+        batch = SingleImageDataset(root_dir='', num_views=6, img_wh=[512, 512], bg_color='white', 
+            crop_size=crop_size, single_image=image, prompt_embeds_path='apps/third_party/Era3D/data/fixed_prompt_embeds_6view')[0]
+        imgs_in = torch.cat([batch['imgs_in']]*2, dim=0)
+        imgs_in = rearrange(imgs_in, "B Nv C H W -> (B Nv) C H W")# (B*Nv, 3, H, W)
+        
+        normal_prompt_embeddings, clr_prompt_embeddings = batch['normal_prompt_embeddings'], batch['color_prompt_embeddings'] 
+        prompt_embeddings = torch.cat([normal_prompt_embeddings, clr_prompt_embeddings], dim=0)
+        prompt_embeddings = rearrange(prompt_embeddings, "B Nv N C -> (B Nv) N C")
+        
+        imgs_in = imgs_in.to(device=device, dtype=torch.float16)
+        prompt_embeddings = prompt_embeddings.to(device=device, dtype=torch.float16)
+        
+        mv_imgs = era3d_pipeline(
+            imgs_in, 
+            None, 
+            prompt_embeds=prompt_embeddings,
+            generator=generator,
+            guidance_scale=guidance_scale, 
+            num_inference_steps=step, 
+            num_images_per_prompt=1, 
+            **{'eta': 1.0}
+        ).images
+        return mv_imgs[6], mv_imgs[8], mv_imgs[9], mv_imgs[10]
 
 @spaces.GPU
 def image2mesh(view_front: np.ndarray, 
@@ -209,24 +242,46 @@ if __name__=="__main__":
         "Auto Remove Background": "Auto Remove Background",
         "Original Image": "Original Image",        
     })
-    mvimg_model_config_list = ["CRM", "ImageDream"]
-    crm_pipeline = TwoStagePipeline(
-                        stage1_model_config,
-                        stage1_sampler_config,
-                        device=device,
-                        dtype=torch.float16
-                    )
-    imagedream_pipeline = MVDreamPipeline.from_pretrained(
-        "ashawkey/imagedream-ipmv-diffusers", # remote weights
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    )
+    mvimg_model_config_list = [
+        "Era3D", 
+        # "CRM", 
+        # "ImageDream"
+    ]
+    if "Era3D" in mvimg_model_config_list:
+        # cfg = load_config("apps/third_party/Era3D/configs/test_unclip-512-6view.yaml")
+        # schema = OmegaConf.structured(TestConfig)
+        # cfg = OmegaConf.merge(schema, cfg)
+        era3d_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
+            'pengHTYX/MacLab-Era3D-512-6view',
+            torch_dtype=torch.float16
+        )
+        # enable xformers
+        era3d_pipeline.unet.enable_xformers_memory_efficient_attention()
+        era3d_pipeline.to(device)
+    elif "CRM" in mvimg_model_config_list:
+        stage1_config = OmegaConf.load(f"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"apps/third_party/CRM/" + stage1_model_config.config
+        crm_pipeline = TwoStagePipeline(
+                            stage1_model_config,
+                            stage1_sampler_config,
+                            device=device,
+                            dtype=torch.float16
+                        )
+    elif "ImageDream" in mvimg_model_config_list:
+        imagedream_pipeline = MVDreamPipeline.from_pretrained(
+            "ashawkey/imagedream-ipmv-diffusers", # remote weights
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+        )
+
+    generator = torch.Generator(device)
 
     # for 3D latent set diffusion
-    ckpt_path = "./ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model.ckpt"
-    config_path = "./ckpts/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml"
-    # 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")
-    # config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml", repo_type="model")
+    ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/model.ckpt", repo_type="model")
+    config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/config.yaml", repo_type="model")
     # ckpt_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/model-300k.ckpt", repo_type="model")
     # config_path = hf_hub_download(repo_id="wyysf/CraftsMan", filename="image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6/config.yaml", repo_type="model")
     scheluder_dict = OrderedDict({ 
@@ -266,7 +321,7 @@ if __name__=="__main__":
                     gr.Markdown('''Try a different <b>seed and MV Model</b> for better results. Good Luck :)''')
                 with gr.Row():
                     seed = gr.Number(0, label='Seed', show_label=True)
-                    mvimg_model = gr.Dropdown(value="CRM", label="MV Image Model", choices=list(mvimg_model_config_list))
+                    mvimg_model = gr.Dropdown(value="Era3D", label="MV Image Model", choices=list(mvimg_model_config_list))
                     more = gr.CheckboxGroup(["Remesh", "Symmetry(TBD)"], label="More", show_label=False)
                 with gr.Row():
                     # input prompt