Upload gradio_app.py

gradio_app.py

@@ -12,9 +12,6 @@ from huggingface_hub import hf_hub_download
 
 from collections import OrderedDict
 import trimesh
-from einops import repeat, rearrange
-import pytorch_lightning as pl
-from typing import Dict, Optional, Tuple, List
 import gradio as gr
 from typing import Any
 
@@ -22,12 +19,8 @@ proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.append(os.path.join(proj_dir))
 
 import tempfile
-import craftsman
-from craftsman.systems.base import BaseSystem
-from craftsman.utils.config import ExperimentConfig, load_config
 
 from apps.utils import *
-from apps.mv_models import GenMVImage
 
 _TITLE = '''CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner'''
 _DESCRIPTION = '''
@@ -64,6 +57,8 @@ CraftsMan is under [AGPL-3.0](https://www.gnu.org/licenses/agpl-3.0.en.html), so
 If you have any questions, feel free to open a discussion or contact us at <b>[email protected]</b>.
 """
 from apps.third_party.CRM.pipelines import TwoStagePipeline
+from apps.third_party.LGM.pipeline_mvdream import MVDreamPipeline
+
 
 model = None
 cached_dir = None
@@ -74,17 +69,43 @@ stage1_model_config.resume = hf_hub_download(repo_id="Zhengyi/CRM", filename="pi
 stage1_model_config.config = f"{parent_dir}/apps/third_party/CRM/" + stage1_model_config.config
 crm_pipeline = None
 
+sys.path.append(f"apps/third_party/LGM")
+imgaedream_pipeline = None
+
+generator = None
+
 @spaces.GPU
 def gen_mvimg(
-    mvimg_model, text, image, crop_size, seed, guidance_scale, step
+    mvimg_model, image, seed, guidance_scale, step, text, neg_text, elevation, 
 ):
-    global crm_pipeline
     if seed == 0:
         seed = np.random.randint(1, 65535)
-    crm_pipeline.set_seed(seed)
-    rt_dict = crm_pipeline(image, scale=guidance_scale, step=step)
-    mv_imgs = rt_dict["stage1_images"]
-    return mv_imgs[5], mv_imgs[3], mv_imgs[2], mv_imgs[0]
+
+    if mvimg_model == "CRM":
+        global crm_pipeline
+        crm_pipeline.set_seed(seed)
+        mv_imgs = crm_pipeline(
+            image, 
+            scale=guidance_scale, 
+            step=step
+        )["stage1_images"]
+        return mv_imgs[5], mv_imgs[3], mv_imgs[2], mv_imgs[0]
+    
+    elif mvimg_model == "ImageDream":
+        global imagedream_pipeline, generator
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4])
+        mv_imgs = imagedream_pipeline(
+            text, 
+            image, 
+            negative_prompt=neg_text, 
+            guidance_scale=guidance_scale,  
+            num_inference_steps=step, 
+            elevation=elevation,
+            generator=generator.manual_seed(seed),
+        )
+        return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
+    
 
 @spaces.GPU
 def image2mesh(view_front: np.ndarray, 
@@ -153,22 +174,27 @@ if __name__=="__main__":
     device = torch.device(f"cuda:{args.device}" if torch.cuda.is_available() else "cpu")
     print(f"using device: {device}")
 
+    # for multi-view images generation
+    background_choice = OrderedDict({ 
+        "Alpha as Mask": "Alpha as Mask",
+        "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,
+    )
+    generator = torch.Generator(device)
+
 
-    # for multi-view images generation
-    background_choice = OrderedDict({ 
-        "Alpha as Mask": "Alpha as Mask",
-        "Auto Remove Background": "Auto Remove Background",
-        "Original Image": "Original Image",        
-    })
-    mvimg_model_config_list = ["CRM"]
-    # mvimg_model_config_list = ["CRM", "ImageDream", "Wonder3D"]
-    
     # 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"
@@ -196,24 +222,33 @@ if __name__=="__main__":
 
         with gr.Row():
             with gr.Column(scale=2):
+                with gr.Column():
+                    # input image
+                    with gr.Row():
+                        image_input = gr.Image(
+                            label="Image Input",
+                            image_mode="RGBA",
+                            sources="upload",
+                            type="pil",
+                        )
+                run_btn = gr.Button('Generate', variant='primary', interactive=True)
+
                 with gr.Row():
-                    image_input = gr.Image(
-                        label="Image Input",
-                        image_mode="RGBA",
-                        sources="upload",
-                        type="pil",
-                    )
-                with gr.Row():
-                    text = gr.Textbox(label="Prompt (Optional, only works for mvdream)", visible=False)
-                with gr.Row():
-                    gr.Markdown('''Try a different <b>seed</b> if the result is unsatisfying. Good Luck :)''')
+                    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))
                     more = gr.CheckboxGroup(["Remesh", "Symmetry(TBD)"], label="More", show_label=False)
-                    # remesh = gr.Checkbox(value=False, label='Remesh')
-                    # symmetry = gr.Checkbox(value=False, label='Symmetry(TBD)', interactive=False)
-                run_btn = gr.Button('Generate', variant='primary', interactive=True)
-                
+                with gr.Row():
+                    # input prompt
+                    text = gr.Textbox(label="Prompt (Opt.)", info="only works for ImageDream")
+
+                with gr.Accordion('Advanced options', open=False):
+                    # negative prompt
+                    neg_text = gr.Textbox(label="Negative Prompt", value='ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate')
+                    # elevation
+                    elevation = gr.Slider(label="elevation", minimum=-90, maximum=90, step=1, value=0)
+
                 with gr.Row():
                     gr.Examples(
                         examples=[os.path.join("./apps/examples", i) for i in os.listdir("./apps/examples")],
@@ -243,10 +278,6 @@ if __name__=="__main__":
                         run_3d_btn = gr.Button('Only Generate 3D', interactive=True)
 
                 with gr.Accordion('Advanced options (2D)', open=False):
-                    with gr.Row():
-                        crop_size = gr.Number(224, label='Crop size')
-                        mvimg_model = gr.Dropdown(value="CRM", label="MV Image Model", choices=mvimg_model_config_list)
-                
                     with gr.Row():
                         foreground_ratio = gr.Slider(
                                 label="Foreground Ratio",
@@ -259,11 +290,11 @@ if __name__=="__main__":
                     with gr.Row():
                         background_choice = gr.Dropdown(label="Backgroud Choice", value="Auto Remove Background",choices=list(background_choice.keys()))
                         rmbg_type = gr.Dropdown(label="Backgroud Remove Type", value="rembg",choices=['sam', "rembg"])
-                        # backgroud_color = gr.ColorPicker(label="Background Color", value="#FFFFFF", interactive=True)
-                        backgroud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=True)
+                        backgroud_color = gr.ColorPicker(label="Background Color", value="#FFFFFF", interactive=True)
+                        # backgroud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=True)
                         
                     with gr.Row():
-                        mvimg_guidance_scale = gr.Number(value=3.5, minimum=3, maximum=10, label="2D Guidance Scale")
+                        mvimg_guidance_scale = gr.Number(value=4.0, minimum=3, maximum=10, label="2D Guidance Scale")
                         mvimg_steps = gr.Number(value=30, minimum=20, maximum=100, label="2D Sample Steps")
             
                 with gr.Accordion('Advanced options (3D)', open=False):
@@ -280,17 +311,16 @@ if __name__=="__main__":
         outputs = [output_model_obj]
         rmbg = RMBG(device)
         
-        # gen_mvimg = GenMVImage(device)
         model = load_model(ckpt_path, config_path, device)
 
         run_btn.click(fn=check_input_image, inputs=[image_input]
                     ).success(
                             fn=rmbg.run, 
-                            inputs=[rmbg_type, image_input, crop_size, foreground_ratio, background_choice, backgroud_color],
+                            inputs=[rmbg_type, image_input, foreground_ratio, background_choice, backgroud_color],
                             outputs=[image_input]
                     ).success(
                             fn=gen_mvimg,
-                            inputs=[mvimg_model, text, image_input, crop_size, seed, mvimg_guidance_scale, mvimg_steps],
+                            inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation],
                             outputs=[view_front, view_right, view_back, view_left]
                     ).success(
                             fn=image2mesh, 
@@ -298,7 +328,7 @@ if __name__=="__main__":
                             outputs=outputs, 
                             api_name="generate_img2obj")
         run_mv_btn.click(fn=gen_mvimg,
-                        inputs=[mvimg_model, text, image_input, crop_size, seed, mvimg_guidance_scale, mvimg_steps],
+                        inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation],
                         outputs=[view_front, view_right, view_back, view_left]
         )
         run_3d_btn.click(fn=image2mesh,