cpuai commited on
Commit
0e28ef0
·
verified ·
1 Parent(s): 33e26d4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +819 -819
app.py CHANGED
@@ -1,819 +1,819 @@
1
- import gradio as gr
2
- from gradio_client import Client, handle_file
3
- import spaces
4
- from concurrent.futures import ThreadPoolExecutor
5
-
6
- import os
7
- os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
8
- os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
9
- os.environ["ATTN_BACKEND"] = "flash_attn_3"
10
- os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
11
- os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
12
- from datetime import datetime
13
- import shutil
14
- import cv2
15
- from typing import *
16
- import torch
17
- import numpy as np
18
- from PIL import Image
19
- import base64
20
- import io
21
- import tempfile
22
- from trellis2.modules.sparse import SparseTensor
23
- from trellis2.pipelines import Trellis2ImageTo3DPipeline
24
- from trellis2.renderers import EnvMap
25
- from trellis2.utils import render_utils
26
- import o_voxel
27
-
28
- # Patch postprocess module with local fix for cumesh.fill_holes() bug
29
- import importlib.util
30
- _local_postprocess = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'o-voxel', 'o_voxel', 'postprocess.py')
31
- if os.path.exists(_local_postprocess):
32
- import sys
33
- _spec = importlib.util.spec_from_file_location('o_voxel.postprocess', _local_postprocess)
34
- _mod = importlib.util.module_from_spec(_spec)
35
- _spec.loader.exec_module(_mod)
36
- o_voxel.postprocess = _mod
37
- sys.modules['o_voxel.postprocess'] = _mod
38
-
39
-
40
- MAX_SEED = np.iinfo(np.int32).max
41
- TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
42
- MODES = [
43
- {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
44
- {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
45
- {"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
46
- {"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
47
- {"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
48
- {"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
49
- ]
50
- STEPS = 8
51
- DEFAULT_MODE = 3
52
- DEFAULT_STEP = 3
53
-
54
-
55
- css = """
56
- /* Overwrite Gradio Default Style */
57
- .stepper-wrapper {
58
- padding: 0;
59
- }
60
-
61
- .stepper-container {
62
- padding: 0;
63
- align-items: center;
64
- }
65
-
66
- .step-button {
67
- flex-direction: row;
68
- }
69
-
70
- .step-connector {
71
- transform: none;
72
- }
73
-
74
- .step-number {
75
- width: 16px;
76
- height: 16px;
77
- }
78
-
79
- .step-label {
80
- position: relative;
81
- bottom: 0;
82
- }
83
-
84
- .wrap.center.full {
85
- inset: 0;
86
- height: 100%;
87
- }
88
-
89
- .wrap.center.full.translucent {
90
- background: var(--block-background-fill);
91
- }
92
-
93
- .meta-text-center {
94
- display: block !important;
95
- position: absolute !important;
96
- top: unset !important;
97
- bottom: 0 !important;
98
- right: 0 !important;
99
- transform: unset !important;
100
- }
101
-
102
- /* Previewer */
103
- .previewer-container {
104
- position: relative;
105
- font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
106
- width: 100%;
107
- height: 722px;
108
- margin: 0 auto;
109
- padding: 20px;
110
- display: flex;
111
- flex-direction: column;
112
- align-items: center;
113
- justify-content: center;
114
- }
115
-
116
- .previewer-container .tips-icon {
117
- position: absolute;
118
- right: 10px;
119
- top: 10px;
120
- z-index: 10;
121
- border-radius: 10px;
122
- color: #fff;
123
- background-color: var(--color-accent);
124
- padding: 3px 6px;
125
- user-select: none;
126
- }
127
-
128
- .previewer-container .tips-text {
129
- position: absolute;
130
- right: 10px;
131
- top: 50px;
132
- color: #fff;
133
- background-color: var(--color-accent);
134
- border-radius: 10px;
135
- padding: 6px;
136
- text-align: left;
137
- max-width: 300px;
138
- z-index: 10;
139
- transition: all 0.3s;
140
- opacity: 0%;
141
- user-select: none;
142
- }
143
-
144
- .previewer-container .tips-text p {
145
- font-size: 14px;
146
- line-height: 1.2;
147
- }
148
-
149
- .tips-icon:hover + .tips-text {
150
- display: block;
151
- opacity: 100%;
152
- }
153
-
154
- /* Row 1: Display Modes */
155
- .previewer-container .mode-row {
156
- width: 100%;
157
- display: flex;
158
- gap: 8px;
159
- justify-content: center;
160
- margin-bottom: 20px;
161
- flex-wrap: wrap;
162
- }
163
- .previewer-container .mode-btn {
164
- width: 24px;
165
- height: 24px;
166
- border-radius: 50%;
167
- cursor: pointer;
168
- opacity: 0.5;
169
- transition: all 0.2s;
170
- border: 2px solid var(--neutral-600, #555);
171
- object-fit: cover;
172
- }
173
- .previewer-container .mode-btn:hover { opacity: 0.9; transform: scale(1.1); }
174
- .previewer-container .mode-btn.active {
175
- opacity: 1;
176
- border-color: var(--color-accent);
177
- transform: scale(1.1);
178
- }
179
-
180
- /* Row 2: Display Image */
181
- .previewer-container .display-row {
182
- margin-bottom: 20px;
183
- min-height: 400px;
184
- width: 100%;
185
- flex-grow: 1;
186
- display: flex;
187
- justify-content: center;
188
- align-items: center;
189
- }
190
- .previewer-container .previewer-main-image {
191
- max-width: 100%;
192
- max-height: 100%;
193
- flex-grow: 1;
194
- object-fit: contain;
195
- display: none;
196
- }
197
- .previewer-container .previewer-main-image.visible {
198
- display: block;
199
- }
200
-
201
- /* Row 3: Custom HTML Slider */
202
- .previewer-container .slider-row {
203
- width: 100%;
204
- display: flex;
205
- flex-direction: column;
206
- align-items: center;
207
- gap: 10px;
208
- padding: 0 10px;
209
- }
210
-
211
- .previewer-container input[type=range] {
212
- -webkit-appearance: none;
213
- width: 100%;
214
- max-width: 400px;
215
- background: transparent;
216
- }
217
- .previewer-container input[type=range]::-webkit-slider-runnable-track {
218
- width: 100%;
219
- height: 8px;
220
- cursor: pointer;
221
- background: var(--neutral-700, #404040);
222
- border-radius: 5px;
223
- }
224
- .previewer-container input[type=range]::-webkit-slider-thumb {
225
- height: 20px;
226
- width: 20px;
227
- border-radius: 50%;
228
- background: var(--color-accent);
229
- cursor: pointer;
230
- -webkit-appearance: none;
231
- margin-top: -6px;
232
- box-shadow: 0 2px 5px rgba(0,0,0,0.2);
233
- transition: transform 0.1s;
234
- }
235
- .previewer-container input[type=range]::-webkit-slider-thumb:hover {
236
- transform: scale(1.2);
237
- }
238
-
239
- /* Overwrite Previewer Block Style */
240
- .gradio-container .padded:has(.previewer-container) {
241
- padding: 0 !important;
242
- }
243
-
244
- .gradio-container:has(.previewer-container) [data-testid="block-label"] {
245
- position: absolute;
246
- top: 0;
247
- left: 0;
248
- }
249
- """
250
-
251
-
252
- head = """
253
- <script>
254
- function refreshView(mode, step) {
255
- // 1. Find current mode and step
256
- const allImgs = document.querySelectorAll('.previewer-main-image');
257
- for (let i = 0; i < allImgs.length; i++) {
258
- const img = allImgs[i];
259
- if (img.classList.contains('visible')) {
260
- const id = img.id;
261
- const [_, m, s] = id.split('-');
262
- if (mode === -1) mode = parseInt(m.slice(1));
263
- if (step === -1) step = parseInt(s.slice(1));
264
- break;
265
- }
266
- }
267
-
268
- // 2. Hide ALL images
269
- // We select all elements with class 'previewer-main-image'
270
- allImgs.forEach(img => img.classList.remove('visible'));
271
-
272
- // 3. Construct the specific ID for the current state
273
- // Format: view-m{mode}-s{step}
274
- const targetId = 'view-m' + mode + '-s' + step;
275
- const targetImg = document.getElementById(targetId);
276
-
277
- // 4. Show ONLY the target
278
- if (targetImg) {
279
- targetImg.classList.add('visible');
280
- }
281
-
282
- // 5. Update Button Highlights
283
- const allBtns = document.querySelectorAll('.mode-btn');
284
- allBtns.forEach((btn, idx) => {
285
- if (idx === mode) btn.classList.add('active');
286
- else btn.classList.remove('active');
287
- });
288
- }
289
-
290
- // --- Action: Switch Mode ---
291
- function selectMode(mode) {
292
- refreshView(mode, -1);
293
- }
294
-
295
- // --- Action: Slider Change ---
296
- function onSliderChange(val) {
297
- refreshView(-1, parseInt(val));
298
- }
299
- </script>
300
- """
301
-
302
-
303
- empty_html = f"""
304
- <div class="previewer-container">
305
- <svg style=" opacity: .5; height: var(--size-5); color: var(--body-text-color);"
306
- xmlns="http://www.w3.org/2000/svg" width="100%" height="100%" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather feather-image"><rect x="3" y="3" width="18" height="18" rx="2" ry="2"></rect><circle cx="8.5" cy="8.5" r="1.5"></circle><polyline points="21 15 16 10 5 21"></polyline></svg>
307
- </div>
308
- """
309
-
310
-
311
- def image_to_base64(image):
312
- buffered = io.BytesIO()
313
- image = image.convert("RGB")
314
- image.save(buffered, format="jpeg", quality=85)
315
- img_str = base64.b64encode(buffered.getvalue()).decode()
316
- return f"data:image/jpeg;base64,{img_str}"
317
-
318
-
319
- def start_session(req: gr.Request):
320
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
321
- os.makedirs(user_dir, exist_ok=True)
322
-
323
-
324
- def end_session(req: gr.Request):
325
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
326
- if os.path.exists(user_dir):
327
- shutil.rmtree(user_dir)
328
-
329
-
330
- def remove_background(input: Image.Image) -> Image.Image:
331
- try:
332
- with tempfile.NamedTemporaryFile(suffix='.png') as f:
333
- input = input.convert('RGB')
334
- input.save(f.name)
335
- output = rmbg_client.predict(handle_file(f.name), api_name="/image")[0][0]
336
- output = Image.open(output)
337
- return output
338
- except Exception as e:
339
- raise gr.Error(f"Background removal failed: {e}. Please upload images with transparent backgrounds (RGBA), or try again later.")
340
-
341
-
342
- def preprocess_image(input: Image.Image) -> Image.Image:
343
- """
344
- Preprocess the input image.
345
- """
346
- # if has alpha channel, use it directly; otherwise, remove background
347
- has_alpha = False
348
- if input.mode == 'RGBA':
349
- alpha = np.array(input)[:, :, 3]
350
- if not np.all(alpha == 255):
351
- has_alpha = True
352
- max_size = max(input.size)
353
- scale = min(1, 1024 / max_size)
354
- if scale < 1:
355
- input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
356
- if has_alpha:
357
- output = input
358
- else:
359
- output = remove_background(input)
360
- output_np = np.array(output)
361
- alpha = output_np[:, :, 3]
362
- bbox = np.argwhere(alpha > 0.8 * 255)
363
- if bbox.size == 0:
364
- # No visible pixels, center the image in a square
365
- size = max(output.size)
366
- square = Image.new('RGB', (size, size), (0, 0, 0))
367
- output_rgb = output.convert('RGB') if output.mode == 'RGBA' else output
368
- square.paste(output_rgb, ((size - output.width) // 2, (size - output.height) // 2))
369
- return square
370
- bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
371
- center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
372
- size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
373
- size = int(size * 1)
374
- bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
375
- output = output.crop(bbox) # type: ignore
376
- output_np = np.array(output).astype(np.float32)
377
- rgb = output_np[:, :, :3]
378
- alpha = output_np[:, :, 3:4] / 255.0
379
- # Keep full RGB for visible pixels, zero out transparent background
380
- mask = (alpha > 0.05).astype(np.float32)
381
- rgb = rgb * mask
382
- output = Image.fromarray(rgb.astype(np.uint8))
383
- return output
384
-
385
-
386
- def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
387
- shape_slat, tex_slat, res = latents
388
- return {
389
- 'shape_slat_feats': shape_slat.feats.cpu().numpy(),
390
- 'tex_slat_feats': tex_slat.feats.cpu().numpy(),
391
- 'coords': shape_slat.coords.cpu().numpy(),
392
- 'res': res,
393
- }
394
-
395
-
396
- def unpack_state(state: dict) -> Tuple[SparseTensor, SparseTensor, int]:
397
- shape_slat = SparseTensor(
398
- feats=torch.from_numpy(state['shape_slat_feats']).cuda(),
399
- coords=torch.from_numpy(state['coords']).cuda(),
400
- )
401
- tex_slat = shape_slat.replace(torch.from_numpy(state['tex_slat_feats']).cuda())
402
- return shape_slat, tex_slat, state['res']
403
-
404
-
405
- def get_seed(randomize_seed, seed):
406
- """
407
- Get the random seed.
408
- """
409
- return np.random.randint(0, MAX_SEED) if randomize_seed else seed
410
-
411
-
412
- def prepare_multi_example() -> List[str]:
413
- """
414
- Prepare multi-image examples. Returns list of image paths.
415
- Shows only the first view as representative thumbnail.
416
- """
417
- multi_case = sorted(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
418
- examples = []
419
- for case in multi_case:
420
- first_img = f'assets/example_multi_image/{case}_1.png'
421
- if os.path.exists(first_img):
422
- examples.append(first_img)
423
- return examples
424
-
425
-
426
- def load_multi_example(image) -> List[Image.Image]:
427
- """Load all views for a multi-image case by matching the input image."""
428
- if image is None:
429
- return []
430
-
431
- # Convert to PIL Image if needed
432
- if isinstance(image, np.ndarray):
433
- image = Image.fromarray(image)
434
-
435
- # Convert to RGB for consistent comparison
436
- input_rgb = np.array(image.convert('RGB'))
437
-
438
- # Find matching case by comparing with first images
439
- example_dir = "assets/example_multi_image"
440
- case_names = sorted(set([f.rsplit('_', 1)[0] for f in os.listdir(example_dir) if f.endswith('.png')]))
441
-
442
- for case_name in case_names:
443
- first_img_path = f'{example_dir}/{case_name}_1.png'
444
- if os.path.exists(first_img_path):
445
- first_img = Image.open(first_img_path).convert('RGB')
446
- first_rgb = np.array(first_img)
447
-
448
- # Compare images (check if same shape and content)
449
- if input_rgb.shape == first_rgb.shape and np.array_equal(input_rgb, first_rgb):
450
- # Found match, load all views (without preprocessing - will be done on Generate)
451
- images = []
452
- for i in range(1, 7):
453
- img_path = f'{example_dir}/{case_name}_{i}.png'
454
- if os.path.exists(img_path):
455
- img = Image.open(img_path).convert('RGBA')
456
- images.append(img)
457
- if images:
458
- return images
459
-
460
- # No match found, return the single image
461
- return [image.convert('RGBA') if image.mode != 'RGBA' else image]
462
-
463
-
464
- def split_image(image: Image.Image) -> List[Image.Image]:
465
- """
466
- Split a concatenated image into multiple views.
467
- """
468
- image = np.array(image)
469
- alpha = image[..., 3]
470
- alpha = np.any(alpha > 0, axis=0)
471
- start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
472
- end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
473
- images = []
474
- for s, e in zip(start_pos, end_pos):
475
- images.append(Image.fromarray(image[:, s:e+1]))
476
- return [preprocess_image(image) for image in images]
477
-
478
-
479
- @spaces.GPU(duration=120)
480
- def image_to_3d(
481
- multiimages,
482
- seed,
483
- resolution,
484
- ss_guidance_strength,
485
- ss_guidance_rescale,
486
- ss_sampling_steps,
487
- ss_rescale_t,
488
- shape_slat_guidance_strength,
489
- shape_slat_guidance_rescale,
490
- shape_slat_sampling_steps,
491
- shape_slat_rescale_t,
492
- tex_slat_guidance_strength,
493
- tex_slat_guidance_rescale,
494
- tex_slat_sampling_steps,
495
- tex_slat_rescale_t,
496
- multiimage_algo,
497
- tex_multiimage_algo,
498
- req: gr.Request,
499
- progress=gr.Progress(track_tqdm=True),
500
- ):
501
- if not multiimages:
502
- raise gr.Error("Please upload images or select an example first.")
503
-
504
- # Preprocess images (background removal for images without alpha)
505
- images = [image[0] for image in multiimages]
506
- processed_images = [preprocess_image(img) for img in images]
507
-
508
- # --- Sampling ---
509
- outputs, latents = pipeline.run_multi_image(
510
- processed_images,
511
- seed=seed,
512
- preprocess_image=False,
513
- sparse_structure_sampler_params={
514
- "steps": ss_sampling_steps,
515
- "guidance_strength": ss_guidance_strength,
516
- "guidance_rescale": ss_guidance_rescale,
517
- "rescale_t": ss_rescale_t,
518
- },
519
- shape_slat_sampler_params={
520
- "steps": shape_slat_sampling_steps,
521
- "guidance_strength": shape_slat_guidance_strength,
522
- "guidance_rescale": shape_slat_guidance_rescale,
523
- "rescale_t": shape_slat_rescale_t,
524
- },
525
- tex_slat_sampler_params={
526
- "steps": tex_slat_sampling_steps,
527
- "guidance_strength": tex_slat_guidance_strength,
528
- "guidance_rescale": tex_slat_guidance_rescale,
529
- "rescale_t": tex_slat_rescale_t,
530
- },
531
- pipeline_type={
532
- "512": "512",
533
- "1024": "1024_cascade",
534
- "1536": "1536_cascade",
535
- }[resolution],
536
- return_latent=True,
537
- mode=multiimage_algo,
538
- tex_mode=tex_multiimage_algo,
539
- )
540
- mesh = outputs[0]
541
- mesh.simplify(16777216) # nvdiffrast limit
542
- images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
543
- state = pack_state(latents)
544
- torch.cuda.empty_cache()
545
-
546
- # --- HTML Construction ---
547
- def encode_preview_image(args):
548
- m_idx, s_idx, render_key = args
549
- img_base64 = image_to_base64(Image.fromarray(images[render_key][s_idx]))
550
- return (m_idx, s_idx, img_base64)
551
-
552
- encode_tasks = [
553
- (m_idx, s_idx, mode['render_key'])
554
- for m_idx, mode in enumerate(MODES)
555
- for s_idx in range(STEPS)
556
- ]
557
-
558
- with ThreadPoolExecutor(max_workers=8) as executor:
559
- encoded_results = list(executor.map(encode_preview_image, encode_tasks))
560
-
561
- encoded_map = {(m, s): b64 for m, s, b64 in encoded_results}
562
- images_html = ""
563
- for m_idx, mode in enumerate(MODES):
564
- for s_idx in range(STEPS):
565
- unique_id = f"view-m{m_idx}-s{s_idx}"
566
- is_visible = (m_idx == DEFAULT_MODE and s_idx == DEFAULT_STEP)
567
- vis_class = "visible" if is_visible else ""
568
- img_base64 = encoded_map[(m_idx, s_idx)]
569
-
570
- images_html += f"""
571
- <img id="{unique_id}"
572
- class="previewer-main-image {vis_class}"
573
- src="{img_base64}"
574
- loading="eager">
575
- """
576
-
577
- btns_html = ""
578
- for idx, mode in enumerate(MODES):
579
- active_class = "active" if idx == DEFAULT_MODE else ""
580
- btns_html += f"""
581
- <img src="{mode['icon_base64']}"
582
- class="mode-btn {active_class}"
583
- onclick="selectMode({idx})"
584
- title="{mode['name']}">
585
- """
586
-
587
- full_html = f"""
588
- <div class="previewer-container">
589
- <div class="tips-wrapper">
590
- <div class="tips-icon">💡Tips</div>
591
- <div class="tips-text">
592
- <p>● <b>Render Mode</b> - Click on the circular buttons to switch between different render modes.</p>
593
- <p>● <b>View Angle</b> - Drag the slider to change the view angle.</p>
594
- </div>
595
- </div>
596
-
597
- <!-- Row 1: Viewport containing 48 static <img> tags -->
598
- <div class="display-row">
599
- {images_html}
600
- </div>
601
-
602
- <!-- Row 2 -->
603
- <div class="mode-row" id="btn-group">
604
- {btns_html}
605
- </div>
606
-
607
- <!-- Row 3: Slider -->
608
- <div class="slider-row">
609
- <input type="range" id="custom-slider" min="0" max="{STEPS - 1}" value="{DEFAULT_STEP}" step="1" oninput="onSliderChange(this.value)">
610
- </div>
611
- </div>
612
- """
613
-
614
- return state, full_html
615
-
616
-
617
- @spaces.GPU(duration=120)
618
- def extract_glb(
619
- state,
620
- decimation_target,
621
- texture_size,
622
- req: gr.Request,
623
- progress=gr.Progress(track_tqdm=True),
624
- ):
625
- """
626
- Extract a GLB file from the 3D model.
627
-
628
- Args:
629
- state (dict): The state of the generated 3D model.
630
- decimation_target (int): The target face count for decimation.
631
- texture_size (int): The texture resolution.
632
-
633
- Returns:
634
- Tuple[str, str]: The path to the extracted GLB file (for Model3D and DownloadButton).
635
- """
636
- user_dir = os.path.join(TMP_DIR, str(req.session_hash))
637
- shape_slat, tex_slat, res = unpack_state(state)
638
- mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
639
- mesh.simplify(16777216) # nvdiffrast limit
640
- glb = o_voxel.postprocess.to_glb(
641
- vertices=mesh.vertices,
642
- faces=mesh.faces,
643
- attr_volume=mesh.attrs,
644
- coords=mesh.coords,
645
- attr_layout=pipeline.pbr_attr_layout,
646
- grid_size=res,
647
- aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
648
- decimation_target=decimation_target,
649
- texture_size=texture_size,
650
- remesh=True,
651
- remesh_band=1,
652
- remesh_project=0,
653
- use_tqdm=True,
654
- )
655
- now = datetime.now()
656
- timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
657
- os.makedirs(user_dir, exist_ok=True)
658
- glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
659
- glb.export(glb_path, extension_webp=True)
660
- torch.cuda.empty_cache()
661
- return glb_path, glb_path
662
-
663
-
664
- with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", neutral_hue="slate"), css=css, head=head) as demo:
665
- gr.HTML("""
666
- <div style="display: flex; align-items: center; gap: 24px;">
667
- <a href="https://www.opsiclear.com" target="_blank" style="flex-shrink: 0; display: flex; align-items: center;">
668
- <img src="https://www.opsiclear.com/assets/logos/Logo_v2_compact_name.svg" alt="OpsiClear"
669
- style="width: 140px; height: auto; object-fit: contain;">
670
- </a>
671
- <div style="min-width: 0; border-left: 2px solid var(--border-color-primary); padding-left: 24px;">
672
- <h2 style="margin: 0 0 8px 0; font-size: 1.4rem; line-height: 1.3; font-weight: 700;">Multi-View to 3D with <a href="https://microsoft.github.io/TRELLIS.2" target="_blank" style="text-decoration: none; color: var(--color-accent);">TRELLIS.2</a></h2>
673
- <ul style="margin: 0; padding-left: 18px; font-size: 0.88rem; line-height: 1.7; color: var(--body-text-color-subdued, var(--body-text-color));">
674
- <li>Upload multiple images from different viewpoints to create a 3D asset with multi-image conditioning.</li>
675
- <li>Click an example below to load a pre-made multi-view set, or upload your own images.</li>
676
- <li>Click <b>Generate</b> to create the 3D model, then <b>Extract GLB</b> to export.</li>
677
- <li style="color: #e67300;"><b>Note:</b> Generation quality is highly sensitive to parameters. Adjust settings in Advanced Settings if results are unsatisfactory.</li>
678
- <li style="color: #cc3333;"><b>Non-Commercial:</b> This space uses models with licenses that <b>forbid commercial use</b> (BRIA RMBG-2.0: CC BY-NC 4.0, nvdiffrast/nvdiffrec: NVIDIA Source Code License).</li>
679
- </ul>
680
- </div>
681
- </div>
682
- """)
683
-
684
- with gr.Row():
685
- with gr.Column(scale=1, min_width=360):
686
- multiimage_prompt = gr.Gallery(label="Multi-View Images", format="png", type="pil", height=400, columns=3, interactive=True)
687
- remove_img_btn = gr.Button("Remove Selected Image", size="sm", variant="secondary")
688
-
689
- resolution = gr.Radio(["512", "1024", "1536"], label="Resolution", value="1024")
690
- seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
691
- randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
692
- decimation_target = gr.Slider(100000, 500000, label="Decimation Target", value=300000, step=10000)
693
- texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
694
-
695
- with gr.Accordion(label="Advanced Settings", open=False):
696
- gr.Markdown("Stage 1: Sparse Structure Generation")
697
- with gr.Row():
698
- ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
699
- ss_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.7, step=0.01)
700
- ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
701
- ss_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=5.0, step=0.1)
702
- gr.Markdown("Stage 2: Shape Generation")
703
- with gr.Row():
704
- shape_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
705
- shape_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.5, step=0.01)
706
- shape_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
707
- shape_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
708
- gr.Markdown("Stage 3: Material Generation")
709
- with gr.Row():
710
- tex_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=1.0, step=0.1)
711
- tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.0, step=0.01)
712
- tex_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
713
- tex_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
714
- multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Structure Algorithm", value="stochastic")
715
- tex_multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Texture Algorithm", value="multidiffusion")
716
-
717
- with gr.Column(scale=10):
718
- preview_output = gr.HTML(empty_html, label="3D Asset Preview", show_label=True, container=True)
719
-
720
- with gr.Row():
721
- generate_btn = gr.Button("Generate", variant="primary")
722
- extract_btn = gr.Button("Extract GLB")
723
-
724
- glb_output = gr.Model3D(label="Extracted GLB", height=600, show_label=True, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0))
725
- download_btn = gr.DownloadButton(label="Download GLB")
726
-
727
- with gr.Accordion(label="Examples", open=True):
728
- example_image = gr.Image(visible=False) # Hidden component for examples
729
- examples_multi = gr.Examples(
730
- examples=prepare_multi_example(),
731
- inputs=[example_image],
732
- fn=load_multi_example,
733
- outputs=[multiimage_prompt],
734
- run_on_click=True,
735
- cache_examples=False,
736
- examples_per_page=50,
737
- )
738
-
739
- output_buf = gr.State()
740
- selected_img_idx = gr.State(value=None)
741
-
742
-
743
- # Handlers
744
- demo.load(start_session)
745
- demo.unload(end_session)
746
-
747
- def on_gallery_select(evt: gr.SelectData):
748
- return evt.index
749
-
750
- def remove_selected_image(images, idx):
751
- if images is None or idx is None or not images:
752
- return images, None
753
- images = list(images)
754
- if idx < len(images):
755
- images.pop(idx)
756
- return images, None
757
-
758
- multiimage_prompt.select(on_gallery_select, outputs=[selected_img_idx])
759
- remove_img_btn.click(
760
- remove_selected_image,
761
- inputs=[multiimage_prompt, selected_img_idx],
762
- outputs=[multiimage_prompt, selected_img_idx],
763
- )
764
-
765
- generate_btn.click(
766
- get_seed,
767
- inputs=[randomize_seed, seed],
768
- outputs=[seed],
769
- ).then(
770
- image_to_3d,
771
- inputs=[
772
- multiimage_prompt, seed, resolution,
773
- ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
774
- shape_slat_guidance_strength, shape_slat_guidance_rescale, shape_slat_sampling_steps, shape_slat_rescale_t,
775
- tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
776
- multiimage_algo, tex_multiimage_algo
777
- ],
778
- outputs=[output_buf, preview_output],
779
- )
780
-
781
- extract_btn.click(
782
- extract_glb,
783
- inputs=[output_buf, decimation_target, texture_size],
784
- outputs=[glb_output, download_btn],
785
- )
786
-
787
-
788
- # Launch the Gradio app
789
- if __name__ == "__main__":
790
- os.makedirs(TMP_DIR, exist_ok=True)
791
-
792
- # Construct ui components
793
- btn_img_base64_strs = {}
794
- for i in range(len(MODES)):
795
- icon = Image.open(MODES[i]['icon'])
796
- MODES[i]['icon_base64'] = image_to_base64(icon)
797
-
798
- rmbg_client = Client("briaai/BRIA-RMBG-2.0")
799
- pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
800
- pipeline.rembg_model = None
801
- pipeline.low_vram = False
802
- pipeline.cuda()
803
-
804
- envmap = {
805
- 'forest': EnvMap(torch.tensor(
806
- cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
807
- dtype=torch.float32, device='cuda'
808
- )),
809
- 'sunset': EnvMap(torch.tensor(
810
- cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
811
- dtype=torch.float32, device='cuda'
812
- )),
813
- 'courtyard': EnvMap(torch.tensor(
814
- cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
815
- dtype=torch.float32, device='cuda'
816
- )),
817
- }
818
-
819
- demo.launch()
 
1
+ import gradio as gr
2
+ from gradio_client import Client, handle_file
3
+ import spaces
4
+ from concurrent.futures import ThreadPoolExecutor
5
+
6
+ import os
7
+ os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
8
+ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
9
+ os.environ["ATTN_BACKEND"] = "flash_attn_3"
10
+ os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
11
+ os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
12
+ from datetime import datetime
13
+ import shutil
14
+ import cv2
15
+ from typing import *
16
+ import torch
17
+ import numpy as np
18
+ from PIL import Image
19
+ import base64
20
+ import io
21
+ import tempfile
22
+ from trellis2.modules.sparse import SparseTensor
23
+ from trellis2.pipelines import Trellis2ImageTo3DPipeline
24
+ from trellis2.renderers import EnvMap
25
+ from trellis2.utils import render_utils
26
+ import o_voxel
27
+
28
+ # Patch postprocess module with local fix for cumesh.fill_holes() bug
29
+ import importlib.util
30
+ _local_postprocess = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'o-voxel', 'o_voxel', 'postprocess.py')
31
+ if os.path.exists(_local_postprocess):
32
+ import sys
33
+ _spec = importlib.util.spec_from_file_location('o_voxel.postprocess', _local_postprocess)
34
+ _mod = importlib.util.module_from_spec(_spec)
35
+ _spec.loader.exec_module(_mod)
36
+ o_voxel.postprocess = _mod
37
+ sys.modules['o_voxel.postprocess'] = _mod
38
+
39
+
40
+ MAX_SEED = np.iinfo(np.int32).max
41
+ TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
42
+ MODES = [
43
+ {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
44
+ {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
45
+ {"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
46
+ {"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
47
+ {"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
48
+ {"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
49
+ ]
50
+ STEPS = 8
51
+ DEFAULT_MODE = 3
52
+ DEFAULT_STEP = 3
53
+
54
+
55
+ css = """
56
+ /* Overwrite Gradio Default Style */
57
+ .stepper-wrapper {
58
+ padding: 0;
59
+ }
60
+
61
+ .stepper-container {
62
+ padding: 0;
63
+ align-items: center;
64
+ }
65
+
66
+ .step-button {
67
+ flex-direction: row;
68
+ }
69
+
70
+ .step-connector {
71
+ transform: none;
72
+ }
73
+
74
+ .step-number {
75
+ width: 16px;
76
+ height: 16px;
77
+ }
78
+
79
+ .step-label {
80
+ position: relative;
81
+ bottom: 0;
82
+ }
83
+
84
+ .wrap.center.full {
85
+ inset: 0;
86
+ height: 100%;
87
+ }
88
+
89
+ .wrap.center.full.translucent {
90
+ background: var(--block-background-fill);
91
+ }
92
+
93
+ .meta-text-center {
94
+ display: block !important;
95
+ position: absolute !important;
96
+ top: unset !important;
97
+ bottom: 0 !important;
98
+ right: 0 !important;
99
+ transform: unset !important;
100
+ }
101
+
102
+ /* Previewer */
103
+ .previewer-container {
104
+ position: relative;
105
+ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
106
+ width: 100%;
107
+ height: 722px;
108
+ margin: 0 auto;
109
+ padding: 20px;
110
+ display: flex;
111
+ flex-direction: column;
112
+ align-items: center;
113
+ justify-content: center;
114
+ }
115
+
116
+ .previewer-container .tips-icon {
117
+ position: absolute;
118
+ right: 10px;
119
+ top: 10px;
120
+ z-index: 10;
121
+ border-radius: 10px;
122
+ color: #fff;
123
+ background-color: var(--color-accent);
124
+ padding: 3px 6px;
125
+ user-select: none;
126
+ }
127
+
128
+ .previewer-container .tips-text {
129
+ position: absolute;
130
+ right: 10px;
131
+ top: 50px;
132
+ color: #fff;
133
+ background-color: var(--color-accent);
134
+ border-radius: 10px;
135
+ padding: 6px;
136
+ text-align: left;
137
+ max-width: 300px;
138
+ z-index: 10;
139
+ transition: all 0.3s;
140
+ opacity: 0%;
141
+ user-select: none;
142
+ }
143
+
144
+ .previewer-container .tips-text p {
145
+ font-size: 14px;
146
+ line-height: 1.2;
147
+ }
148
+
149
+ .tips-icon:hover + .tips-text {
150
+ display: block;
151
+ opacity: 100%;
152
+ }
153
+
154
+ /* Row 1: Display Modes */
155
+ .previewer-container .mode-row {
156
+ width: 100%;
157
+ display: flex;
158
+ gap: 8px;
159
+ justify-content: center;
160
+ margin-bottom: 20px;
161
+ flex-wrap: wrap;
162
+ }
163
+ .previewer-container .mode-btn {
164
+ width: 24px;
165
+ height: 24px;
166
+ border-radius: 50%;
167
+ cursor: pointer;
168
+ opacity: 0.5;
169
+ transition: all 0.2s;
170
+ border: 2px solid var(--neutral-600, #555);
171
+ object-fit: cover;
172
+ }
173
+ .previewer-container .mode-btn:hover { opacity: 0.9; transform: scale(1.1); }
174
+ .previewer-container .mode-btn.active {
175
+ opacity: 1;
176
+ border-color: var(--color-accent);
177
+ transform: scale(1.1);
178
+ }
179
+
180
+ /* Row 2: Display Image */
181
+ .previewer-container .display-row {
182
+ margin-bottom: 20px;
183
+ min-height: 400px;
184
+ width: 100%;
185
+ flex-grow: 1;
186
+ display: flex;
187
+ justify-content: center;
188
+ align-items: center;
189
+ }
190
+ .previewer-container .previewer-main-image {
191
+ max-width: 100%;
192
+ max-height: 100%;
193
+ flex-grow: 1;
194
+ object-fit: contain;
195
+ display: none;
196
+ }
197
+ .previewer-container .previewer-main-image.visible {
198
+ display: block;
199
+ }
200
+
201
+ /* Row 3: Custom HTML Slider */
202
+ .previewer-container .slider-row {
203
+ width: 100%;
204
+ display: flex;
205
+ flex-direction: column;
206
+ align-items: center;
207
+ gap: 10px;
208
+ padding: 0 10px;
209
+ }
210
+
211
+ .previewer-container input[type=range] {
212
+ -webkit-appearance: none;
213
+ width: 100%;
214
+ max-width: 400px;
215
+ background: transparent;
216
+ }
217
+ .previewer-container input[type=range]::-webkit-slider-runnable-track {
218
+ width: 100%;
219
+ height: 8px;
220
+ cursor: pointer;
221
+ background: var(--neutral-700, #404040);
222
+ border-radius: 5px;
223
+ }
224
+ .previewer-container input[type=range]::-webkit-slider-thumb {
225
+ height: 20px;
226
+ width: 20px;
227
+ border-radius: 50%;
228
+ background: var(--color-accent);
229
+ cursor: pointer;
230
+ -webkit-appearance: none;
231
+ margin-top: -6px;
232
+ box-shadow: 0 2px 5px rgba(0,0,0,0.2);
233
+ transition: transform 0.1s;
234
+ }
235
+ .previewer-container input[type=range]::-webkit-slider-thumb:hover {
236
+ transform: scale(1.2);
237
+ }
238
+
239
+ /* Overwrite Previewer Block Style */
240
+ .gradio-container .padded:has(.previewer-container) {
241
+ padding: 0 !important;
242
+ }
243
+
244
+ .gradio-container:has(.previewer-container) [data-testid="block-label"] {
245
+ position: absolute;
246
+ top: 0;
247
+ left: 0;
248
+ }
249
+ """
250
+
251
+
252
+ head = """
253
+ <script>
254
+ function refreshView(mode, step) {
255
+ // 1. Find current mode and step
256
+ const allImgs = document.querySelectorAll('.previewer-main-image');
257
+ for (let i = 0; i < allImgs.length; i++) {
258
+ const img = allImgs[i];
259
+ if (img.classList.contains('visible')) {
260
+ const id = img.id;
261
+ const [_, m, s] = id.split('-');
262
+ if (mode === -1) mode = parseInt(m.slice(1));
263
+ if (step === -1) step = parseInt(s.slice(1));
264
+ break;
265
+ }
266
+ }
267
+
268
+ // 2. Hide ALL images
269
+ // We select all elements with class 'previewer-main-image'
270
+ allImgs.forEach(img => img.classList.remove('visible'));
271
+
272
+ // 3. Construct the specific ID for the current state
273
+ // Format: view-m{mode}-s{step}
274
+ const targetId = 'view-m' + mode + '-s' + step;
275
+ const targetImg = document.getElementById(targetId);
276
+
277
+ // 4. Show ONLY the target
278
+ if (targetImg) {
279
+ targetImg.classList.add('visible');
280
+ }
281
+
282
+ // 5. Update Button Highlights
283
+ const allBtns = document.querySelectorAll('.mode-btn');
284
+ allBtns.forEach((btn, idx) => {
285
+ if (idx === mode) btn.classList.add('active');
286
+ else btn.classList.remove('active');
287
+ });
288
+ }
289
+
290
+ // --- Action: Switch Mode ---
291
+ function selectMode(mode) {
292
+ refreshView(mode, -1);
293
+ }
294
+
295
+ // --- Action: Slider Change ---
296
+ function onSliderChange(val) {
297
+ refreshView(-1, parseInt(val));
298
+ }
299
+ </script>
300
+ """
301
+
302
+
303
+ empty_html = f"""
304
+ <div class="previewer-container">
305
+ <svg style=" opacity: .5; height: var(--size-5); color: var(--body-text-color);"
306
+ xmlns="http://www.w3.org/2000/svg" width="100%" height="100%" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather feather-image"><rect x="3" y="3" width="18" height="18" rx="2" ry="2"></rect><circle cx="8.5" cy="8.5" r="1.5"></circle><polyline points="21 15 16 10 5 21"></polyline></svg>
307
+ </div>
308
+ """
309
+
310
+
311
+ def image_to_base64(image):
312
+ buffered = io.BytesIO()
313
+ image = image.convert("RGB")
314
+ image.save(buffered, format="jpeg", quality=85)
315
+ img_str = base64.b64encode(buffered.getvalue()).decode()
316
+ return f"data:image/jpeg;base64,{img_str}"
317
+
318
+
319
+ def start_session(req: gr.Request):
320
+ user_dir = os.path.join(TMP_DIR, str(req.session_hash))
321
+ os.makedirs(user_dir, exist_ok=True)
322
+
323
+
324
+ def end_session(req: gr.Request):
325
+ user_dir = os.path.join(TMP_DIR, str(req.session_hash))
326
+ if os.path.exists(user_dir):
327
+ shutil.rmtree(user_dir)
328
+
329
+
330
+ def remove_background(input: Image.Image) -> Image.Image:
331
+ try:
332
+ with tempfile.NamedTemporaryFile(suffix='.png') as f:
333
+ input = input.convert('RGB')
334
+ input.save(f.name)
335
+ output = rmbg_client.predict(handle_file(f.name), api_name="/image")[0][0]
336
+ output = Image.open(output)
337
+ return output
338
+ except Exception as e:
339
+ raise gr.Error(f"Background removal failed: {e}. Please upload images with transparent backgrounds (RGBA), or try again later.")
340
+
341
+
342
+ def preprocess_image(input: Image.Image) -> Image.Image:
343
+ """
344
+ Preprocess the input image.
345
+ """
346
+ # if has alpha channel, use it directly; otherwise, remove background
347
+ has_alpha = False
348
+ if input.mode == 'RGBA':
349
+ alpha = np.array(input)[:, :, 3]
350
+ if not np.all(alpha == 255):
351
+ has_alpha = True
352
+ max_size = max(input.size)
353
+ scale = min(1, 1024 / max_size)
354
+ if scale < 1:
355
+ input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
356
+ if has_alpha:
357
+ output = input
358
+ else:
359
+ output = remove_background(input)
360
+ output_np = np.array(output)
361
+ alpha = output_np[:, :, 3]
362
+ bbox = np.argwhere(alpha > 0.8 * 255)
363
+ if bbox.size == 0:
364
+ # No visible pixels, center the image in a square
365
+ size = max(output.size)
366
+ square = Image.new('RGB', (size, size), (0, 0, 0))
367
+ output_rgb = output.convert('RGB') if output.mode == 'RGBA' else output
368
+ square.paste(output_rgb, ((size - output.width) // 2, (size - output.height) // 2))
369
+ return square
370
+ bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
371
+ center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
372
+ size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
373
+ size = int(size * 1)
374
+ bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
375
+ output = output.crop(bbox) # type: ignore
376
+ output_np = np.array(output).astype(np.float32)
377
+ rgb = output_np[:, :, :3]
378
+ alpha = output_np[:, :, 3:4] / 255.0
379
+ # Keep full RGB for visible pixels, zero out transparent background
380
+ mask = (alpha > 0.05).astype(np.float32)
381
+ rgb = rgb * mask
382
+ output = Image.fromarray(rgb.astype(np.uint8))
383
+ return output
384
+
385
+
386
+ def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
387
+ shape_slat, tex_slat, res = latents
388
+ return {
389
+ 'shape_slat_feats': shape_slat.feats.cpu().numpy(),
390
+ 'tex_slat_feats': tex_slat.feats.cpu().numpy(),
391
+ 'coords': shape_slat.coords.cpu().numpy(),
392
+ 'res': res,
393
+ }
394
+
395
+
396
+ def unpack_state(state: dict) -> Tuple[SparseTensor, SparseTensor, int]:
397
+ shape_slat = SparseTensor(
398
+ feats=torch.from_numpy(state['shape_slat_feats']).cuda(),
399
+ coords=torch.from_numpy(state['coords']).cuda(),
400
+ )
401
+ tex_slat = shape_slat.replace(torch.from_numpy(state['tex_slat_feats']).cuda())
402
+ return shape_slat, tex_slat, state['res']
403
+
404
+
405
+ def get_seed(randomize_seed, seed):
406
+ """
407
+ Get the random seed.
408
+ """
409
+ return np.random.randint(0, MAX_SEED) if randomize_seed else seed
410
+
411
+
412
+ def prepare_multi_example() -> List[str]:
413
+ """
414
+ Prepare multi-image examples. Returns list of image paths.
415
+ Shows only the first view as representative thumbnail.
416
+ """
417
+ multi_case = sorted(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
418
+ examples = []
419
+ for case in multi_case:
420
+ first_img = f'assets/example_multi_image/{case}_1.png'
421
+ if os.path.exists(first_img):
422
+ examples.append(first_img)
423
+ return examples
424
+
425
+
426
+ def load_multi_example(image) -> List[Image.Image]:
427
+ """Load all views for a multi-image case by matching the input image."""
428
+ if image is None:
429
+ return []
430
+
431
+ # Convert to PIL Image if needed
432
+ if isinstance(image, np.ndarray):
433
+ image = Image.fromarray(image)
434
+
435
+ # Convert to RGB for consistent comparison
436
+ input_rgb = np.array(image.convert('RGB'))
437
+
438
+ # Find matching case by comparing with first images
439
+ example_dir = "assets/example_multi_image"
440
+ case_names = sorted(set([f.rsplit('_', 1)[0] for f in os.listdir(example_dir) if f.endswith('.png')]))
441
+
442
+ for case_name in case_names:
443
+ first_img_path = f'{example_dir}/{case_name}_1.png'
444
+ if os.path.exists(first_img_path):
445
+ first_img = Image.open(first_img_path).convert('RGB')
446
+ first_rgb = np.array(first_img)
447
+
448
+ # Compare images (check if same shape and content)
449
+ if input_rgb.shape == first_rgb.shape and np.array_equal(input_rgb, first_rgb):
450
+ # Found match, load all views (without preprocessing - will be done on Generate)
451
+ images = []
452
+ for i in range(1, 7):
453
+ img_path = f'{example_dir}/{case_name}_{i}.png'
454
+ if os.path.exists(img_path):
455
+ img = Image.open(img_path).convert('RGBA')
456
+ images.append(img)
457
+ if images:
458
+ return images
459
+
460
+ # No match found, return the single image
461
+ return [image.convert('RGBA') if image.mode != 'RGBA' else image]
462
+
463
+
464
+ def split_image(image: Image.Image) -> List[Image.Image]:
465
+ """
466
+ Split a concatenated image into multiple views.
467
+ """
468
+ image = np.array(image)
469
+ alpha = image[..., 3]
470
+ alpha = np.any(alpha > 0, axis=0)
471
+ start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
472
+ end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
473
+ images = []
474
+ for s, e in zip(start_pos, end_pos):
475
+ images.append(Image.fromarray(image[:, s:e+1]))
476
+ return [preprocess_image(image) for image in images]
477
+
478
+
479
+ @spaces.GPU(duration=120)
480
+ def image_to_3d(
481
+ multiimages,
482
+ seed,
483
+ resolution,
484
+ ss_guidance_strength,
485
+ ss_guidance_rescale,
486
+ ss_sampling_steps,
487
+ ss_rescale_t,
488
+ shape_slat_guidance_strength,
489
+ shape_slat_guidance_rescale,
490
+ shape_slat_sampling_steps,
491
+ shape_slat_rescale_t,
492
+ tex_slat_guidance_strength,
493
+ tex_slat_guidance_rescale,
494
+ tex_slat_sampling_steps,
495
+ tex_slat_rescale_t,
496
+ multiimage_algo,
497
+ tex_multiimage_algo,
498
+ req: gr.Request,
499
+ progress=gr.Progress(track_tqdm=True),
500
+ ):
501
+ if not multiimages:
502
+ raise gr.Error("Please upload images or select an example first.")
503
+
504
+ # Preprocess images (background removal for images without alpha)
505
+ images = [image[0] for image in multiimages]
506
+ processed_images = [preprocess_image(img) for img in images]
507
+
508
+ # --- Sampling ---
509
+ outputs, latents = pipeline.run_multi_image(
510
+ processed_images,
511
+ seed=seed,
512
+ preprocess_image=False,
513
+ sparse_structure_sampler_params={
514
+ "steps": ss_sampling_steps,
515
+ "guidance_strength": ss_guidance_strength,
516
+ "guidance_rescale": ss_guidance_rescale,
517
+ "rescale_t": ss_rescale_t,
518
+ },
519
+ shape_slat_sampler_params={
520
+ "steps": shape_slat_sampling_steps,
521
+ "guidance_strength": shape_slat_guidance_strength,
522
+ "guidance_rescale": shape_slat_guidance_rescale,
523
+ "rescale_t": shape_slat_rescale_t,
524
+ },
525
+ tex_slat_sampler_params={
526
+ "steps": tex_slat_sampling_steps,
527
+ "guidance_strength": tex_slat_guidance_strength,
528
+ "guidance_rescale": tex_slat_guidance_rescale,
529
+ "rescale_t": tex_slat_rescale_t,
530
+ },
531
+ pipeline_type={
532
+ "512": "512",
533
+ "1024": "1024_cascade",
534
+ "1536": "1536_cascade",
535
+ }[resolution],
536
+ return_latent=True,
537
+ mode=multiimage_algo,
538
+ tex_mode=tex_multiimage_algo,
539
+ )
540
+ mesh = outputs[0]
541
+ mesh.simplify(16777216) # nvdiffrast limit
542
+ images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
543
+ state = pack_state(latents)
544
+ torch.cuda.empty_cache()
545
+
546
+ # --- HTML Construction ---
547
+ def encode_preview_image(args):
548
+ m_idx, s_idx, render_key = args
549
+ img_base64 = image_to_base64(Image.fromarray(images[render_key][s_idx]))
550
+ return (m_idx, s_idx, img_base64)
551
+
552
+ encode_tasks = [
553
+ (m_idx, s_idx, mode['render_key'])
554
+ for m_idx, mode in enumerate(MODES)
555
+ for s_idx in range(STEPS)
556
+ ]
557
+
558
+ with ThreadPoolExecutor(max_workers=8) as executor:
559
+ encoded_results = list(executor.map(encode_preview_image, encode_tasks))
560
+
561
+ encoded_map = {(m, s): b64 for m, s, b64 in encoded_results}
562
+ images_html = ""
563
+ for m_idx, mode in enumerate(MODES):
564
+ for s_idx in range(STEPS):
565
+ unique_id = f"view-m{m_idx}-s{s_idx}"
566
+ is_visible = (m_idx == DEFAULT_MODE and s_idx == DEFAULT_STEP)
567
+ vis_class = "visible" if is_visible else ""
568
+ img_base64 = encoded_map[(m_idx, s_idx)]
569
+
570
+ images_html += f"""
571
+ <img id="{unique_id}"
572
+ class="previewer-main-image {vis_class}"
573
+ src="{img_base64}"
574
+ loading="eager">
575
+ """
576
+
577
+ btns_html = ""
578
+ for idx, mode in enumerate(MODES):
579
+ active_class = "active" if idx == DEFAULT_MODE else ""
580
+ btns_html += f"""
581
+ <img src="{mode['icon_base64']}"
582
+ class="mode-btn {active_class}"
583
+ onclick="selectMode({idx})"
584
+ title="{mode['name']}">
585
+ """
586
+
587
+ full_html = f"""
588
+ <div class="previewer-container">
589
+ <div class="tips-wrapper">
590
+ <div class="tips-icon">💡Tips</div>
591
+ <div class="tips-text">
592
+ <p>● <b>Render Mode</b> - Click on the circular buttons to switch between different render modes.</p>
593
+ <p>● <b>View Angle</b> - Drag the slider to change the view angle.</p>
594
+ </div>
595
+ </div>
596
+
597
+ <!-- Row 1: Viewport containing 48 static <img> tags -->
598
+ <div class="display-row">
599
+ {images_html}
600
+ </div>
601
+
602
+ <!-- Row 2 -->
603
+ <div class="mode-row" id="btn-group">
604
+ {btns_html}
605
+ </div>
606
+
607
+ <!-- Row 3: Slider -->
608
+ <div class="slider-row">
609
+ <input type="range" id="custom-slider" min="0" max="{STEPS - 1}" value="{DEFAULT_STEP}" step="1" oninput="onSliderChange(this.value)">
610
+ </div>
611
+ </div>
612
+ """
613
+
614
+ return state, full_html
615
+
616
+
617
+ @spaces.GPU(duration=120)
618
+ def extract_glb(
619
+ state,
620
+ decimation_target,
621
+ texture_size,
622
+ req: gr.Request,
623
+ progress=gr.Progress(track_tqdm=True),
624
+ ):
625
+ """
626
+ Extract a GLB file from the 3D model.
627
+
628
+ Args:
629
+ state (dict): The state of the generated 3D model.
630
+ decimation_target (int): The target face count for decimation.
631
+ texture_size (int): The texture resolution.
632
+
633
+ Returns:
634
+ Tuple[str, str]: The path to the extracted GLB file (for Model3D and DownloadButton).
635
+ """
636
+ user_dir = os.path.join(TMP_DIR, str(req.session_hash))
637
+ shape_slat, tex_slat, res = unpack_state(state)
638
+ mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
639
+ mesh.simplify(16777216) # nvdiffrast limit
640
+ glb = o_voxel.postprocess.to_glb(
641
+ vertices=mesh.vertices,
642
+ faces=mesh.faces,
643
+ attr_volume=mesh.attrs,
644
+ coords=mesh.coords,
645
+ attr_layout=pipeline.pbr_attr_layout,
646
+ grid_size=res,
647
+ aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
648
+ decimation_target=decimation_target,
649
+ texture_size=texture_size,
650
+ remesh=True,
651
+ remesh_band=1,
652
+ remesh_project=0,
653
+ use_tqdm=True,
654
+ )
655
+ now = datetime.now()
656
+ timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
657
+ os.makedirs(user_dir, exist_ok=True)
658
+ glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
659
+ glb.export(glb_path, extension_webp=False)
660
+ torch.cuda.empty_cache()
661
+ return glb_path, glb_path
662
+
663
+
664
+ with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", neutral_hue="slate"), css=css, head=head) as demo:
665
+ gr.HTML("""
666
+ <div style="display: flex; align-items: center; gap: 24px;">
667
+ <a href="https://www.opsiclear.com" target="_blank" style="flex-shrink: 0; display: flex; align-items: center;">
668
+ <img src="https://www.opsiclear.com/assets/logos/Logo_v2_compact_name.svg" alt="OpsiClear"
669
+ style="width: 140px; height: auto; object-fit: contain;">
670
+ </a>
671
+ <div style="min-width: 0; border-left: 2px solid var(--border-color-primary); padding-left: 24px;">
672
+ <h2 style="margin: 0 0 8px 0; font-size: 1.4rem; line-height: 1.3; font-weight: 700;">Multi-View to 3D with <a href="https://microsoft.github.io/TRELLIS.2" target="_blank" style="text-decoration: none; color: var(--color-accent);">TRELLIS.2</a></h2>
673
+ <ul style="margin: 0; padding-left: 18px; font-size: 0.88rem; line-height: 1.7; color: var(--body-text-color-subdued, var(--body-text-color));">
674
+ <li>Upload multiple images from different viewpoints to create a 3D asset with multi-image conditioning.</li>
675
+ <li>Click an example below to load a pre-made multi-view set, or upload your own images.</li>
676
+ <li>Click <b>Generate</b> to create the 3D model, then <b>Extract GLB</b> to export.</li>
677
+ <li style="color: #e67300;"><b>Note:</b> Generation quality is highly sensitive to parameters. Adjust settings in Advanced Settings if results are unsatisfactory.</li>
678
+ <li style="color: #cc3333;"><b>Non-Commercial:</b> This space uses models with licenses that <b>forbid commercial use</b> (BRIA RMBG-2.0: CC BY-NC 4.0, nvdiffrast/nvdiffrec: NVIDIA Source Code License).</li>
679
+ </ul>
680
+ </div>
681
+ </div>
682
+ """)
683
+
684
+ with gr.Row():
685
+ with gr.Column(scale=1, min_width=360):
686
+ multiimage_prompt = gr.Gallery(label="Multi-View Images", format="png", type="pil", height=400, columns=3, interactive=True)
687
+ remove_img_btn = gr.Button("Remove Selected Image", size="sm", variant="secondary")
688
+
689
+ resolution = gr.Radio(["512", "1024", "1536"], label="Resolution", value="1024")
690
+ seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
691
+ randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
692
+ decimation_target = gr.Slider(100000, 500000, label="Decimation Target", value=300000, step=10000)
693
+ texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
694
+
695
+ with gr.Accordion(label="Advanced Settings", open=False):
696
+ gr.Markdown("Stage 1: Sparse Structure Generation")
697
+ with gr.Row():
698
+ ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
699
+ ss_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.7, step=0.01)
700
+ ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
701
+ ss_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=5.0, step=0.1)
702
+ gr.Markdown("Stage 2: Shape Generation")
703
+ with gr.Row():
704
+ shape_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
705
+ shape_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.5, step=0.01)
706
+ shape_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
707
+ shape_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
708
+ gr.Markdown("Stage 3: Material Generation")
709
+ with gr.Row():
710
+ tex_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=1.0, step=0.1)
711
+ tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.0, step=0.01)
712
+ tex_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
713
+ tex_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
714
+ multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Structure Algorithm", value="stochastic")
715
+ tex_multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Texture Algorithm", value="multidiffusion")
716
+
717
+ with gr.Column(scale=10):
718
+ preview_output = gr.HTML(empty_html, label="3D Asset Preview", show_label=True, container=True)
719
+
720
+ with gr.Row():
721
+ generate_btn = gr.Button("Generate", variant="primary")
722
+ extract_btn = gr.Button("Extract GLB")
723
+
724
+ glb_output = gr.Model3D(label="Extracted GLB", height=600, show_label=True, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0))
725
+ download_btn = gr.DownloadButton(label="Download GLB")
726
+
727
+ with gr.Accordion(label="Examples", open=True):
728
+ example_image = gr.Image(visible=False) # Hidden component for examples
729
+ examples_multi = gr.Examples(
730
+ examples=prepare_multi_example(),
731
+ inputs=[example_image],
732
+ fn=load_multi_example,
733
+ outputs=[multiimage_prompt],
734
+ run_on_click=True,
735
+ cache_examples=False,
736
+ examples_per_page=50,
737
+ )
738
+
739
+ output_buf = gr.State()
740
+ selected_img_idx = gr.State(value=None)
741
+
742
+
743
+ # Handlers
744
+ demo.load(start_session)
745
+ demo.unload(end_session)
746
+
747
+ def on_gallery_select(evt: gr.SelectData):
748
+ return evt.index
749
+
750
+ def remove_selected_image(images, idx):
751
+ if images is None or idx is None or not images:
752
+ return images, None
753
+ images = list(images)
754
+ if idx < len(images):
755
+ images.pop(idx)
756
+ return images, None
757
+
758
+ multiimage_prompt.select(on_gallery_select, outputs=[selected_img_idx])
759
+ remove_img_btn.click(
760
+ remove_selected_image,
761
+ inputs=[multiimage_prompt, selected_img_idx],
762
+ outputs=[multiimage_prompt, selected_img_idx],
763
+ )
764
+
765
+ generate_btn.click(
766
+ get_seed,
767
+ inputs=[randomize_seed, seed],
768
+ outputs=[seed],
769
+ ).then(
770
+ image_to_3d,
771
+ inputs=[
772
+ multiimage_prompt, seed, resolution,
773
+ ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
774
+ shape_slat_guidance_strength, shape_slat_guidance_rescale, shape_slat_sampling_steps, shape_slat_rescale_t,
775
+ tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
776
+ multiimage_algo, tex_multiimage_algo
777
+ ],
778
+ outputs=[output_buf, preview_output],
779
+ )
780
+
781
+ extract_btn.click(
782
+ extract_glb,
783
+ inputs=[output_buf, decimation_target, texture_size],
784
+ outputs=[glb_output, download_btn],
785
+ )
786
+
787
+
788
+ # Launch the Gradio app
789
+ if __name__ == "__main__":
790
+ os.makedirs(TMP_DIR, exist_ok=True)
791
+
792
+ # Construct ui components
793
+ btn_img_base64_strs = {}
794
+ for i in range(len(MODES)):
795
+ icon = Image.open(MODES[i]['icon'])
796
+ MODES[i]['icon_base64'] = image_to_base64(icon)
797
+
798
+ rmbg_client = Client("briaai/BRIA-RMBG-2.0")
799
+ pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
800
+ pipeline.rembg_model = None
801
+ pipeline.low_vram = False
802
+ pipeline.cuda()
803
+
804
+ envmap = {
805
+ 'forest': EnvMap(torch.tensor(
806
+ cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
807
+ dtype=torch.float32, device='cuda'
808
+ )),
809
+ 'sunset': EnvMap(torch.tensor(
810
+ cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
811
+ dtype=torch.float32, device='cuda'
812
+ )),
813
+ 'courtyard': EnvMap(torch.tensor(
814
+ cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
815
+ dtype=torch.float32, device='cuda'
816
+ )),
817
+ }
818
+
819
+ demo.launch()