davidvgilmore commited on
Commit
c5efce8
·
verified ·
1 Parent(s): b54ca85

Upload hg_app_bak.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. hg_app_bak.py +402 -0
hg_app_bak.py ADDED
@@ -0,0 +1,402 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # pip install gradio==4.44.1
2
+ if True:
3
+ import os
4
+ import spaces
5
+ import subprocess
6
+ import sys
7
+ import shlex
8
+ print("cd /home/user/app/hy3dgen/texgen/differentiable_renderer/ && bash compile_mesh_painter.sh")
9
+ os.system("cd /home/user/app/hy3dgen/texgen/differentiable_renderer/ && bash compile_mesh_painter.sh")
10
+ print('install custom')
11
+ subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
12
+ IP = "0.0.0.0"
13
+ PORT = 7860
14
+ else:
15
+ IP = "0.0.0.0"
16
+ PORT = 8080
17
+ class spaces:
18
+ class GPU:
19
+ def __init__(self, duration=60):
20
+ self.duration = duration
21
+ def __call__(self, func):
22
+ return func
23
+
24
+ import os
25
+ import shutil
26
+ import time
27
+ from glob import glob
28
+ import gradio as gr
29
+ import torch
30
+ from gradio_litmodel3d import LitModel3D
31
+
32
+
33
+ def get_example_img_list():
34
+ print('Loading example img list ...')
35
+ return sorted(glob('./assets/example_images/*.png'))
36
+
37
+
38
+ def get_example_txt_list():
39
+ print('Loading example txt list ...')
40
+ txt_list = list()
41
+ for line in open('./assets/example_prompts.txt'):
42
+ txt_list.append(line.strip())
43
+ return txt_list
44
+
45
+
46
+ def gen_save_folder(max_size=6000):
47
+ os.makedirs(SAVE_DIR, exist_ok=True)
48
+ exists = set(int(_) for _ in os.listdir(SAVE_DIR) if not _.startswith("."))
49
+ cur_id = min(set(range(max_size)) - exists) if len(exists) < max_size else -1
50
+ if os.path.exists(f"{SAVE_DIR}/{(cur_id + 1) % max_size}"):
51
+ shutil.rmtree(f"{SAVE_DIR}/{(cur_id + 1) % max_size}")
52
+ print(f"remove {SAVE_DIR}/{(cur_id + 1) % max_size} success !!!")
53
+ save_folder = f"{SAVE_DIR}/{max(0, cur_id)}"
54
+ os.makedirs(save_folder, exist_ok=True)
55
+ print(f"mkdir {save_folder} suceess !!!")
56
+ return save_folder
57
+
58
+
59
+ def export_mesh(mesh, save_folder, textured=False):
60
+ if textured:
61
+ path = os.path.join(save_folder, f'textured_mesh.glb')
62
+ else:
63
+ path = os.path.join(save_folder, f'white_mesh.glb')
64
+ mesh.export(path, include_normals=textured)
65
+ return path
66
+
67
+
68
+ def build_model_viewer_html(save_folder, height=660, width=790, textured=False):
69
+ if textured:
70
+ related_path = f"./textured_mesh.glb"
71
+ template_name = './assets/modelviewer-textured-template.html'
72
+ output_html_path = os.path.join(save_folder, f'textured_mesh.html')
73
+ else:
74
+ related_path = f"./white_mesh.glb"
75
+ template_name = './assets/modelviewer-template.html'
76
+ output_html_path = os.path.join(save_folder, f'white_mesh.html')
77
+
78
+ with open(os.path.join(CURRENT_DIR, template_name), 'r') as f:
79
+ template_html = f.read()
80
+ obj_html = f"""
81
+ <div class="column is-mobile is-centered">
82
+ <model-viewer style="height: {height - 10}px; width: {width}px;" rotation-per-second="10deg" id="modelViewer"
83
+ src="{related_path}/" disable-tap
84
+ environment-image="neutral" auto-rotate camera-target="0m 0m 0m" orientation="0deg 0deg 170deg" shadow-intensity=".9"
85
+ ar auto-rotate camera-controls>
86
+ </model-viewer>
87
+ </div>
88
+ """
89
+
90
+ with open(output_html_path, 'w') as f:
91
+ f.write(template_html.replace('<model-viewer>', obj_html))
92
+
93
+ iframe_tag = f'<iframe src="file/{output_html_path}" height="{height}" width="100%" frameborder="0"></iframe>'
94
+ print(f'Find html {output_html_path}, {os.path.exists(output_html_path)}')
95
+
96
+ return f"""
97
+ <div style='height: {height}; width: 100%;'>
98
+ {iframe_tag}
99
+ </div>
100
+ """
101
+
102
+ @spaces.GPU(duration=60)
103
+ def _gen_shape(
104
+ caption,
105
+ image,
106
+ steps=50,
107
+ guidance_scale=7.5,
108
+ seed=1234,
109
+ octree_resolution=256,
110
+ check_box_rembg=False,
111
+ ):
112
+ if caption: print('prompt is', caption)
113
+ save_folder = gen_save_folder()
114
+ stats = {}
115
+ time_meta = {}
116
+ start_time_0 = time.time()
117
+
118
+ image_path = ''
119
+ if image is None:
120
+ start_time = time.time()
121
+ image = t2i_worker(caption)
122
+ time_meta['text2image'] = time.time() - start_time
123
+
124
+ image.save(os.path.join(save_folder, 'input.png'))
125
+
126
+ print(image.mode)
127
+ if check_box_rembg or image.mode == "RGB":
128
+ start_time = time.time()
129
+ image = rmbg_worker(image.convert('RGB'))
130
+ time_meta['rembg'] = time.time() - start_time
131
+
132
+ image.save(os.path.join(save_folder, 'rembg.png'))
133
+
134
+ # image to white model
135
+ start_time = time.time()
136
+
137
+ generator = torch.Generator()
138
+ generator = generator.manual_seed(int(seed))
139
+ mesh = i23d_worker(
140
+ image=image,
141
+ num_inference_steps=steps,
142
+ guidance_scale=guidance_scale,
143
+ generator=generator,
144
+ octree_resolution=octree_resolution
145
+ )[0]
146
+
147
+ mesh = FloaterRemover()(mesh)
148
+ mesh = DegenerateFaceRemover()(mesh)
149
+ mesh = FaceReducer()(mesh)
150
+
151
+ stats['number_of_faces'] = mesh.faces.shape[0]
152
+ stats['number_of_vertices'] = mesh.vertices.shape[0]
153
+
154
+ time_meta['image_to_textured_3d'] = {'total': time.time() - start_time}
155
+ time_meta['total'] = time.time() - start_time_0
156
+ stats['time'] = time_meta
157
+ return mesh, save_folder, image
158
+
159
+ @spaces.GPU(duration=80)
160
+ def generation_all(
161
+ caption,
162
+ image,
163
+ steps=50,
164
+ guidance_scale=7.5,
165
+ seed=1234,
166
+ octree_resolution=256,
167
+ check_box_rembg=False
168
+ ):
169
+ mesh, save_folder, image = _gen_shape(
170
+ caption,
171
+ image,
172
+ steps=steps,
173
+ guidance_scale=guidance_scale,
174
+ seed=seed,
175
+ octree_resolution=octree_resolution,
176
+ check_box_rembg=check_box_rembg
177
+ )
178
+ path = export_mesh(mesh, save_folder, textured=False)
179
+ model_viewer_html = build_model_viewer_html(save_folder, height=596, width=700)
180
+
181
+ textured_mesh = texgen_worker(mesh, image)
182
+ path_textured = export_mesh(textured_mesh, save_folder, textured=True)
183
+ model_viewer_html_textured = build_model_viewer_html(save_folder, height=596, width=700, textured=True)
184
+
185
+ return (
186
+ gr.update(value=path, visible=True),
187
+ gr.update(value=path_textured, visible=True),
188
+ gr.update(value=path, visible=True),
189
+ gr.update(value=path_textured, visible=True),
190
+ # model_viewer_html,
191
+ # model_viewer_html_textured,
192
+ )
193
+
194
+ @spaces.GPU(duration=30)
195
+ def shape_generation(
196
+ caption,
197
+ image,
198
+ steps=50,
199
+ guidance_scale=7.5,
200
+ seed=1234,
201
+ octree_resolution=256,
202
+ check_box_rembg=False,
203
+ ):
204
+ mesh, save_folder, image = _gen_shape(
205
+ caption,
206
+ image,
207
+ steps=steps,
208
+ guidance_scale=guidance_scale,
209
+ seed=seed,
210
+ octree_resolution=octree_resolution,
211
+ check_box_rembg=check_box_rembg
212
+ )
213
+
214
+ path = export_mesh(mesh, save_folder, textured=False)
215
+ model_viewer_html = build_model_viewer_html(save_folder, height=596, width=700)
216
+
217
+ return (
218
+ gr.update(value=path, visible=True),
219
+ gr.update(value=path, visible=True),
220
+ # model_viewer_html,
221
+ )
222
+
223
+
224
+ def build_app():
225
+ title_html = """
226
+ <div style="font-size: 2em; font-weight: bold; text-align: center; margin-bottom: 20px">
227
+
228
+ Hunyuan3D-2: Scaling Diffusion Models for High Resolution Textured 3D Assets Generation
229
+ </div>
230
+ <div align="center">
231
+ Tencent Hunyuan3D Team
232
+ </div>
233
+ <div align="center">
234
+ <a href="https://github.com/tencent/Hunyuan3D-1">Github Page</a> &ensp;
235
+ <a href="http://3d-models.hunyuan.tencent.com">Homepage</a> &ensp;
236
+ <a href="https://arxiv.org/pdf/2411.02293">Technical Report</a> &ensp;
237
+ <a href="https://huggingface.co/Tencent/Hunyuan3D-2"> Models</a> &ensp;
238
+ </div>
239
+ """
240
+ css = """
241
+ .json-output {
242
+ height: 578px;
243
+ }
244
+ .json-output .json-holder {
245
+ height: 538px;
246
+ overflow-y: scroll;
247
+ }
248
+ """
249
+
250
+ with gr.Blocks(theme=gr.themes.Base(), css=css, title='Hunyuan-3D-2.0') as demo:
251
+ # if not gr.__version__.startswith('4'): gr.HTML(title_html)
252
+ gr.HTML(title_html)
253
+
254
+ with gr.Row():
255
+ with gr.Column(scale=2):
256
+ with gr.Tabs() as tabs_prompt:
257
+ with gr.Tab('Image Prompt', id='tab_img_prompt') as tab_ip:
258
+ image = gr.Image(label='Image', type='pil', image_mode='RGBA', height=290)
259
+ with gr.Row():
260
+ check_box_rembg = gr.Checkbox(value=True, label='Remove Background')
261
+
262
+ with gr.Tab('Text Prompt', id='tab_txt_prompt') as tab_tp:
263
+ caption = gr.Textbox(label='Text Prompt',
264
+ placeholder='HunyuanDiT will be used to generate image.',
265
+ info='Example: A 3D model of a cute cat, white background')
266
+
267
+ with gr.Accordion('Advanced Options', open=False):
268
+ num_steps = gr.Slider(maximum=50, minimum=20, value=30, step=1, label='Inference Steps')
269
+ octree_resolution = gr.Dropdown([256, 384, 512], value=256, label='Octree Resolution')
270
+ cfg_scale = gr.Number(value=5.5, label='Guidance Scale')
271
+ seed = gr.Slider(maximum=1e7, minimum=0, value=1234, label='Seed')
272
+
273
+ with gr.Group():
274
+ btn = gr.Button(value='Generate Shape Only', variant='primary')
275
+ btn_all = gr.Button(value='Generate Shape and Texture', variant='primary')
276
+
277
+ with gr.Group():
278
+ file_out = gr.File(label="File", visible=False)
279
+ file_out2 = gr.File(label="File", visible=False)
280
+
281
+ with gr.Column(scale=5):
282
+ with gr.Tabs():
283
+ with gr.Tab('Generated Mesh') as mesh1:
284
+ mesh_output1 = LitModel3D(
285
+ label="3D Model1",
286
+ exposure=10.0,
287
+ height=600,
288
+ visible=True,
289
+ clear_color=[0.0, 0.0, 0.0, 0.0],
290
+ tonemapping="aces",
291
+ contrast=1.0,
292
+ scale=1.0,
293
+ )
294
+ # html_output1 = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output')
295
+ with gr.Tab('Generated Textured Mesh') as mesh2:
296
+ # html_output2 = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output')
297
+ mesh_output2 = LitModel3D(
298
+ label="3D Model2",
299
+ exposure=10.0,
300
+ height=600,
301
+ visible=True,
302
+ clear_color=[0.0, 0.0, 0.0, 0.0],
303
+ tonemapping="aces",
304
+ contrast=1.0,
305
+ scale=1.0,
306
+ )
307
+
308
+ with gr.Column(scale=2):
309
+ with gr.Tabs() as gallery:
310
+ with gr.Tab('Image to 3D Gallery', id='tab_img_gallery') as tab_gi:
311
+ with gr.Row():
312
+ gr.Examples(examples=example_is, inputs=[image],
313
+ label="Image Prompts", examples_per_page=18)
314
+
315
+ with gr.Tab('Text to 3D Gallery', id='tab_txt_gallery') as tab_gt:
316
+ with gr.Row():
317
+ gr.Examples(examples=example_ts, inputs=[caption],
318
+ label="Text Prompts", examples_per_page=18)
319
+
320
+ tab_gi.select(fn=lambda: gr.update(selected='tab_img_prompt'), outputs=tabs_prompt)
321
+ tab_gt.select(fn=lambda: gr.update(selected='tab_txt_prompt'), outputs=tabs_prompt)
322
+
323
+ btn.click(
324
+ shape_generation,
325
+ inputs=[
326
+ caption,
327
+ image,
328
+ num_steps,
329
+ cfg_scale,
330
+ seed,
331
+ octree_resolution,
332
+ check_box_rembg,
333
+ ],
334
+ # outputs=[file_out, html_output1]
335
+ outputs=[file_out, mesh_output1]
336
+ ).then(
337
+ lambda: gr.update(visible=True),
338
+ outputs=[file_out],
339
+ )
340
+
341
+ btn_all.click(
342
+ generation_all,
343
+ inputs=[
344
+ caption,
345
+ image,
346
+ num_steps,
347
+ cfg_scale,
348
+ seed,
349
+ octree_resolution,
350
+ check_box_rembg,
351
+ ],
352
+ # outputs=[file_out, file_out2, html_output1, html_output2]
353
+ outputs=[file_out, file_out2, mesh_output1, mesh_output2]
354
+ ).then(
355
+ lambda: (gr.update(visible=True), gr.update(visible=True)),
356
+ outputs=[file_out, file_out2],
357
+ )
358
+
359
+ return demo
360
+
361
+
362
+ if __name__ == '__main__':
363
+ import argparse
364
+
365
+ parser = argparse.ArgumentParser()
366
+ parser.add_argument('--port', type=int, default=8080)
367
+ parser.add_argument('--cache-path', type=str, default='./gradio_cache')
368
+ args = parser.parse_args()
369
+
370
+ SAVE_DIR = args.cache_path
371
+ os.makedirs(SAVE_DIR, exist_ok=True)
372
+
373
+ CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
374
+
375
+ HTML_OUTPUT_PLACEHOLDER = """
376
+ <div style='height: 596px; width: 100%; border-radius: 8px; border-color: #e5e7eb; order-style: solid; border-width: 1px;'></div>
377
+ """
378
+
379
+ INPUT_MESH_HTML = """
380
+ <div style='height: 490px; width: 100%; border-radius: 8px;
381
+ border-color: #e5e7eb; order-style: solid; border-width: 1px;'>
382
+ </div>
383
+ """
384
+ example_is = get_example_img_list()
385
+ example_ts = get_example_txt_list()
386
+
387
+ from hy3dgen.text2image import HunyuanDiTPipeline
388
+ from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover, \
389
+ Hunyuan3DDiTFlowMatchingPipeline
390
+ from hy3dgen.texgen import Hunyuan3DPaintPipeline
391
+ from hy3dgen.rembg import BackgroundRemover
392
+
393
+ rmbg_worker = BackgroundRemover()
394
+ t2i_worker = HunyuanDiTPipeline()
395
+ i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained('tencent/Hunyuan3D-2')
396
+ texgen_worker = Hunyuan3DPaintPipeline.from_pretrained('tencent/Hunyuan3D-2')
397
+ floater_remove_worker = FloaterRemover()
398
+ degenerate_face_remove_worker = DegenerateFaceRemover()
399
+ face_reduce_worker = FaceReducer()
400
+
401
+ demo = build_app()
402
+ demo.queue().launch(server_name=IP,server_port=PORT)