File size: 17,136 Bytes
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
516314a
302c8a6
516314a
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
df36cab
302c8a6
 
df36cab
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
9dd707f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2711cda
302c8a6
 
 
 
 
 
 
2711cda
 
302c8a6
 
 
 
 
 
 
 
 
2711cda
 
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2b55a44
 
302c8a6
 
2b55a44
 
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
df36cab
302c8a6
 
 
 
10c831f
302c8a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2711cda
302c8a6
 
 
 
 
 
 
2711cda
302c8a6
 
 
 
 
 
 
0f079b2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
import spaces
import argparse
import os
import json
import torch
import sys
import time
import importlib
import numpy as np
from omegaconf import OmegaConf
from huggingface_hub import hf_hub_download

from collections import OrderedDict
import trimesh
import gradio as gr
from typing import Any

proj_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(proj_dir))

import tempfile

from apps.utils import *

_TITLE = '''CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner'''
_DESCRIPTION = '''
<div>
<span style="color: red;">Important: The ckpt models released have been primarily trained on character data, hence they are likely to exhibit superior performance in this category. We are also planning to release more advanced pretrained models in the future.</span>
<br>
By mimicking the artist/craftsman modeling workflow, we propose CraftsMan (aka εŒ εΏƒ) which uses 3D Latent Set Diffusion Model that directly generates coarse meshes,
then a multi-view normal enhanced image generation model is used to refine the mesh.
We provide the coarse 3D diffusion part here. 
<br>
If you found CraftsMan is helpful, please help to ⭐ the <a href='https://github.com/wyysf-98/CraftsMan/' target='_blank'>Github Repo</a>. Thanks!
<a style="display:inline-block; margin-left: .5em" href='https://github.com/wyysf-98/CraftsMan/'><img src='https://img.shields.io/github/stars/wyysf-98/CraftsMan?style=social' /></a>
<br>
*If you have your own multi-view images, you can directly upload it.
</div>
'''
_CITE_ = r"""
---
πŸ“ **Citation**
If you find our work useful for your research or applications, please cite using this bibtex:
```bibtex
@article{li2024craftsman,
author    = {Weiyu Li and Jiarui Liu and Rui Chen and Yixun Liang and Xuelin Chen and Ping Tan and Xiaoxiao Long},
title     = {CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner},
journal   = {arXiv preprint arXiv:2405.14979},
year      = {2024},
}
```
πŸ€— **Acknowledgements**
We use <a href='https://github.com/wjakob/instant-meshes' target='_blank'>Instant Meshes</a> to remesh the generated mesh to a lower face count, thanks to the authors for the great work.
πŸ“‹ **License**
CraftsMan is under [AGPL-3.0](https://www.gnu.org/licenses/agpl-3.0.en.html), so any downstream solution and products (including cloud services) that include CraftsMan code or a trained model (both pretrained or custom trained) inside it should be open-sourced to comply with the AGPL conditions. If you have any questions about the usage of CraftsMan, please contact us first.
πŸ“§ **Contact**
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

import re
import os
import stat

RD, WD, XD = 4, 2, 1
BNS = [RD, WD, XD]
MDS = [
    [stat.S_IRUSR, stat.S_IRGRP, stat.S_IROTH],
    [stat.S_IWUSR, stat.S_IWGRP, stat.S_IWOTH],
    [stat.S_IXUSR, stat.S_IXGRP, stat.S_IXOTH]
]

def chmod(path, mode):
    if isinstance(mode, int):
        mode = str(mode)
    if not re.match("^[0-7]{1,3}$", mode):
        raise Exception("mode does not conform to ^[0-7]{1,3}$ pattern")
    mode = "{0:0>3}".format(mode)
    mode_num = 0
    for midx, m in enumerate(mode):
        for bnidx, bn in enumerate(BNS):
            if (int(m) & bn) > 0:
                mode_num += MDS[bnidx][midx]
    os.chmod(path, mode_num)

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
crm_pipeline = None

sys.path.append(f"apps/third_party/LGM")
imgaedream_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)

    if mvimg_model == "CRM":
        global crm_pipeline
        crm_pipeline.set_seed(seed)
        background = Image.new("RGBA", image.size, (127, 127, 127))
        image = Image.alpha_composite(background, image)
        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
        background = Image.new("RGBA", image.size, backgroud_color)
        image = Image.alpha_composite(background, image)
        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,
        )
        return mv_imgs[1], mv_imgs[2], mv_imgs[3], mv_imgs[0]
    

@spaces.GPU
def image2mesh(view_front: np.ndarray, 
               view_right: np.ndarray, 
               view_back: np.ndarray, 
               view_left: np.ndarray,
               more: bool = False,
               scheluder_name: str ="DDIMScheduler",
               guidance_scale: int = 7.5,
               seed: int = 4,
               octree_depth: int = 7):
    
    sample_inputs = {
        "mvimages": [[
            Image.fromarray(view_front), 
            Image.fromarray(view_right), 
            Image.fromarray(view_back), 
            Image.fromarray(view_left)
        ]]
    }

    global model
    latents = model.sample(
        sample_inputs,
        sample_times=1,
        guidance_scale=guidance_scale,
        return_intermediates=False,
        seed=seed
       
    )[0]
    
    # decode the latents to mesh
    box_v = 1.1
    mesh_outputs, _ = model.shape_model.extract_geometry(
        latents,
        bounds=[-box_v, -box_v, -box_v, box_v, box_v, box_v],
        octree_depth=octree_depth
    )
    assert len(mesh_outputs) == 1, "Only support single mesh output for gradio demo"
    mesh = trimesh.Trimesh(mesh_outputs[0][0], mesh_outputs[0][1])
    # filepath = f"{cached_dir}/{time.time()}.obj"
    filepath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name
    mesh.export(filepath, include_normals=True)

    if 'Remesh' in more:
        remeshed_filepath = tempfile.NamedTemporaryFile(suffix=f"_remeshed.obj", delete=False).name
        print("Remeshing with Instant Meshes...")
        # target_face_count = int(len(mesh.faces)/10)
        target_face_count = 2000
        command = f"{proj_dir}/apps/third_party/InstantMeshes {filepath} -f {target_face_count} -o {remeshed_filepath}"
        os.system(command)
        filepath = remeshed_filepath
        # filepath = filepath.replace('.obj', '_remeshed.obj')
    
    return filepath

if __name__=="__main__":
    parser = argparse.ArgumentParser()
    # parser.add_argument("--model_path", type=str, required=True, help="Path to the object file",)
    parser.add_argument("--cached_dir", type=str, default="./gradio_cached_dir")
    parser.add_argument("--device", type=int, default=0)
    args = parser.parse_args()

    cached_dir = args.cached_dir
    os.makedirs(args.cached_dir, exist_ok=True)
    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,
    )

    # 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/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({ 
        "DDIMScheduler": 'diffusers.schedulers.DDIMScheduler',
        # "DPMSolverMultistepScheduler": 'diffusers.schedulers.DPMSolverMultistepScheduler', # not support yet
        # "UniPCMultistepScheduler": 'diffusers.schedulers.UniPCMultistepScheduler', # not support yet
    })
    
    # main GUI
    custom_theme = gr.themes.Soft(primary_hue="blue").set(
                    button_secondary_background_fill="*neutral_100",
                    button_secondary_background_fill_hover="*neutral_200")
    custom_css = '''#disp_image {
        text-align: center; /* Horizontally center the content */
    }'''
    
    with gr.Blocks(title=_TITLE, theme=custom_theme, css=custom_css) as demo:
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown('# ' + _TITLE)
        gr.Markdown(_DESCRIPTION)

        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():
                    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)
                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")],
                        inputs=[image_input],
                        examples_per_page=8
                    )

            with gr.Column(scale=4):
                with gr.Row():
                    output_model_obj = gr.Model3D(
                        label="Output Model (OBJ Format)",
                        camera_position=(90.0, 90.0, 3.5),
                        interactive=False,
                    )
                with gr.Row():
                    gr.Markdown('''*please note that the model is fliped due to the gradio viewer, please download the obj file and you will get the correct orientation.''')
                    
                with gr.Row():
                    view_front = gr.Image(label="Front", interactive=True, show_label=True)
                    view_right = gr.Image(label="Right", interactive=True, show_label=True)
                    view_back = gr.Image(label="Back", interactive=True, show_label=True)
                    view_left = gr.Image(label="Left", interactive=True, show_label=True)
                    
                with gr.Accordion('Advanced options', open=False):
                    with gr.Row(equal_height=True):
                        run_mv_btn = gr.Button('Only Generate 2D', interactive=True)
                        run_3d_btn = gr.Button('Only Generate 3D', interactive=True)

                with gr.Accordion('Advanced options (2D)', open=False):
                    with gr.Row():
                        foreground_ratio = gr.Slider(
                                label="Foreground Ratio",
                                minimum=0.5,
                                maximum=1.0,
                                value=1.0,
                                step=0.05,
                            )
                        
                    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)
                        
                    with gr.Row():
                        mvimg_guidance_scale = gr.Number(value=3.0, minimum=1, 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):
                    with gr.Row():
                        guidance_scale = gr.Number(label="3D Guidance Scale", value=3.0, minimum=1.0, maximum=10.0)
                        steps = gr.Number(value=50, minimum=20, maximum=100, label="3D Sample Steps")
                        
                    with gr.Row():
                        scheduler = gr.Dropdown(label="scheluder", value="DDIMScheduler",choices=list(scheluder_dict.keys()))
                        octree_depth = gr.Slider(label="Octree Depth", value=7, minimum=4, maximum=8, step=1)
                    
        gr.Markdown(_CITE_)

        outputs = [output_model_obj]
        rmbg = RMBG(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, foreground_ratio, background_choice, backgroud_color],
                            outputs=[image_input]
                    ).success(
                            fn=gen_mvimg,
                            inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation, backgroud_color],
                            outputs=[view_front, view_right, view_back, view_left]
                    ).success(
                            fn=image2mesh, 
                            inputs=[view_front, view_right, view_back, view_left, more, scheduler, guidance_scale, seed, octree_depth],
                            outputs=outputs, 
                            api_name="generate_img2obj")
        run_mv_btn.click(fn=gen_mvimg,
                        inputs=[mvimg_model, image_input, seed, mvimg_guidance_scale, mvimg_steps, text, neg_text, elevation, backgroud_color],
                        outputs=[view_front, view_right, view_back, view_left]
        )
        run_3d_btn.click(fn=image2mesh, 
                        inputs=[view_front, view_right, view_back, view_left, more, scheduler, guidance_scale, seed, octree_depth],
                        outputs=outputs, 
                        api_name="generate_img2obj")
        
        demo.queue().launch(share=True, allowed_paths=[args.cached_dir])