staswrs commited on
Commit
e28dbf7
·
1 Parent(s): ecfd160

add octree depth controls

Browse files
Files changed (3) hide show
  1. app.py +12 -2
  2. app_backlog.py +191 -46
  3. inference_triposg.py +2 -0
app.py CHANGED
@@ -67,7 +67,9 @@ rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
67
  rmbg_net.eval()
68
 
69
 
70
- def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
 
 
71
  print("[API CALL] image_path received:", image_path)
72
  print("[API CALL] File exists:", os.path.exists(image_path))
73
 
@@ -84,6 +86,7 @@ def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
84
  num_inference_steps=int(num_steps),
85
  guidance_scale=float(guidance_scale),
86
  faces=int(face_number),
 
87
  )
88
 
89
  if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
@@ -125,7 +128,14 @@ def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
125
  # Интерфейс Gradio
126
  demo = gr.Interface(
127
  fn=generate,
128
- inputs=gr.Image(type="filepath", label="Upload image"),
 
 
 
 
 
 
 
129
  outputs=gr.File(label="Download .glb"),
130
  title="TripoSG Image to 3D",
131
  description="Upload an image to generate a 3D model (.glb)",
 
67
  rmbg_net.eval()
68
 
69
 
70
+ # def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
71
+ def generate(image_path, face_number=50000, octree_depth=9, guidance_scale=5.0, num_steps=25):
72
+ print(f"[INPUT] octree_depth={octree_depth}, face_number={face_number}, guidance_scale={guidance_scale}, num_steps={num_steps}")# 👈 добавлено
73
  print("[API CALL] image_path received:", image_path)
74
  print("[API CALL] File exists:", os.path.exists(image_path))
75
 
 
86
  num_inference_steps=int(num_steps),
87
  guidance_scale=float(guidance_scale),
88
  faces=int(face_number),
89
+ octree_depth=int(octree_depth), # 👈 добавлено
90
  )
91
 
92
  if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
 
128
  # Интерфейс Gradio
129
  demo = gr.Interface(
130
  fn=generate,
131
+ # inputs=gr.Image(type="filepath", label="Upload image"),
132
+ inputs=[
133
+ gr.Image(type="filepath", label="Upload image"),
134
+ gr.Slider(10000, 150000, step=10000, value=50000, label="Face count"),
135
+ gr.Slider(6, 10, step=1, value=9, label="Octree Depth"),
136
+ gr.Slider(1.0, 10.0, step=0.5, value=5.0, label="Guidance Scale"),
137
+ gr.Slider(10, 100, step=5, value=25, label="Steps"),
138
+ ], # 👈 добавлено
139
  outputs=gr.File(label="Download .glb"),
140
  title="TripoSG Image to 3D",
141
  description="Upload an image to generate a 3D model (.glb)",
app_backlog.py CHANGED
@@ -302,6 +302,192 @@
302
 
303
 
304
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
305
  import os
306
  import subprocess
307
 
@@ -324,12 +510,15 @@ import requests
324
  import traceback
325
  import trimesh
326
  import numpy as np
 
 
327
  from trimesh.exchange.gltf import export_glb
328
 
329
  from inference_triposg import run_triposg
330
  from triposg.pipelines.pipeline_triposg import TripoSGPipeline
331
  from briarmbg import BriaRMBG
332
 
 
333
 
334
  print("Trimesh version:", trimesh.__version__)
335
 
@@ -365,51 +554,6 @@ pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
365
  rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
366
  rmbg_net.eval()
367
 
368
- # Генерация .glb
369
- # def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
370
- # print("[API CALL] image_path received:", image_path)
371
- # print("[API CALL] File exists:", os.path.exists(image_path))
372
-
373
- # temp_id = str(uuid.uuid4())
374
- # output_path = f"/tmp/{temp_id}.glb"
375
- # print("[DEBUG] Generating mesh from:", image_path)
376
-
377
- # try:
378
- # mesh = run_triposg(
379
- # pipe=pipe,
380
- # image_input=image_path,
381
- # rmbg_net=rmbg_net,
382
- # seed=42,
383
- # num_inference_steps=int(num_steps),
384
- # guidance_scale=float(guidance_scale),
385
- # faces=int(face_number),
386
- # )
387
-
388
- # if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
389
- # raise ValueError("Mesh generation returned an empty mesh")
390
-
391
- # mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
392
- # mesh.rezero()
393
- # mesh.fix_normals()
394
- # mesh.apply_translation(-mesh.center_mass)
395
-
396
- # # Масштабируем, чтобы модель вписывалась в размер 1x1x1
397
- # # Если нужно будет подгонять под размер в Endless Tools, то можно использовать:
398
- # # scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
399
- # # mesh.apply_scale(scale_factor)
400
-
401
-
402
- # glb_data = mesh.export(file_type='glb')
403
- # with open(output_path, "wb") as f:
404
- # f.write(glb_data)
405
-
406
- # print(f"[DEBUG] Mesh saved to {output_path}")
407
- # return output_path if os.path.exists(output_path) else None
408
-
409
- # except Exception as e:
410
- # print("[ERROR]", e)
411
- # traceback.print_exc()
412
- # return f"Error: {e}"
413
 
414
  def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
415
  print("[API CALL] image_path received:", image_path)
@@ -452,7 +596,6 @@ def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
452
  else:
453
  print("[DEBUG] Normals missing.")
454
 
455
-
456
  # 💾 Сохраняем GLB
457
  glb_data = mesh.export(file_type='glb')
458
  with open(output_path, "wb") as f:
@@ -476,5 +619,7 @@ demo = gr.Interface(
476
  description="Upload an image to generate a 3D model (.glb)",
477
  )
478
 
 
 
479
  # Запуск
480
  demo.launch()
 
302
 
303
 
304
 
305
+ # import os
306
+ # import subprocess
307
+
308
+ # # Убираем pyenv
309
+ # os.environ.pop("PYENV_VERSION", None)
310
+
311
+ # # Установка зависимостей
312
+ # subprocess.run(["pip", "install", "torch", "wheel"], check=True)
313
+ # subprocess.run([
314
+ # "pip", "install", "--no-build-isolation",
315
+ # "diso@git+https://github.com/SarahWeiii/diso.git"
316
+ # ], check=True)
317
+
318
+ # # Импорты (перенесены после установки зависимостей)
319
+ # import gradio as gr
320
+ # import uuid
321
+ # import torch
322
+ # import zipfile
323
+ # import requests
324
+ # import traceback
325
+ # import trimesh
326
+ # import numpy as np
327
+ # from trimesh.exchange.gltf import export_glb
328
+
329
+ # from inference_triposg import run_triposg
330
+ # from triposg.pipelines.pipeline_triposg import TripoSGPipeline
331
+ # from briarmbg import BriaRMBG
332
+
333
+
334
+ # print("Trimesh version:", trimesh.__version__)
335
+
336
+ # # Настройки устройства
337
+ # device = "cuda" if torch.cuda.is_available() else "cpu"
338
+ # dtype = torch.float16 if device == "cuda" else torch.float32
339
+
340
+ # # Загрузка весов
341
+ # weights_dir = "pretrained_weights"
342
+ # triposg_path = os.path.join(weights_dir, "TripoSG")
343
+ # rmbg_path = os.path.join(weights_dir, "RMBG-1.4")
344
+
345
+ # if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
346
+ # print("📦 Downloading pretrained weights...")
347
+ # url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
348
+ # zip_path = "pretrained_models.zip"
349
+
350
+ # with requests.get(url, stream=True) as r:
351
+ # r.raise_for_status()
352
+ # with open(zip_path, "wb") as f:
353
+ # for chunk in r.iter_content(chunk_size=8192):
354
+ # f.write(chunk)
355
+
356
+ # print("📦 Extracting weights...")
357
+ # with zipfile.ZipFile(zip_path, "r") as zip_ref:
358
+ # zip_ref.extractall(weights_dir)
359
+
360
+ # os.remove(zip_path)
361
+ # print("✅ Weights ready.")
362
+
363
+ # # Загрузка моделей
364
+ # pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
365
+ # rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
366
+ # rmbg_net.eval()
367
+
368
+ # # Генерация .glb
369
+ # # def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
370
+ # # print("[API CALL] image_path received:", image_path)
371
+ # # print("[API CALL] File exists:", os.path.exists(image_path))
372
+
373
+ # # temp_id = str(uuid.uuid4())
374
+ # # output_path = f"/tmp/{temp_id}.glb"
375
+ # # print("[DEBUG] Generating mesh from:", image_path)
376
+
377
+ # # try:
378
+ # # mesh = run_triposg(
379
+ # # pipe=pipe,
380
+ # # image_input=image_path,
381
+ # # rmbg_net=rmbg_net,
382
+ # # seed=42,
383
+ # # num_inference_steps=int(num_steps),
384
+ # # guidance_scale=float(guidance_scale),
385
+ # # faces=int(face_number),
386
+ # # )
387
+
388
+ # # if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
389
+ # # raise ValueError("Mesh generation returned an empty mesh")
390
+
391
+ # # mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
392
+ # # mesh.rezero()
393
+ # # mesh.fix_normals()
394
+ # # mesh.apply_translation(-mesh.center_mass)
395
+
396
+ # # # Масштабируем, чтобы модель вписывалась в размер 1x1x1
397
+ # # # Если нужно будет подгонять под размер в Endless Tools, то можно использовать:
398
+ # # # scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
399
+ # # # mesh.apply_scale(scale_factor)
400
+
401
+
402
+ # # glb_data = mesh.export(file_type='glb')
403
+ # # with open(output_path, "wb") as f:
404
+ # # f.write(glb_data)
405
+
406
+ # # print(f"[DEBUG] Mesh saved to {output_path}")
407
+ # # return output_path if os.path.exists(output_path) else None
408
+
409
+ # # except Exception as e:
410
+ # # print("[ERROR]", e)
411
+ # # traceback.print_exc()
412
+ # # return f"Error: {e}"
413
+
414
+ # def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
415
+ # print("[API CALL] image_path received:", image_path)
416
+ # print("[API CALL] File exists:", os.path.exists(image_path))
417
+
418
+ # temp_id = str(uuid.uuid4())
419
+ # output_path = f"/tmp/{temp_id}.glb"
420
+ # print("[DEBUG] Generating mesh from:", image_path)
421
+
422
+ # try:
423
+ # mesh = run_triposg(
424
+ # pipe=pipe,
425
+ # image_input=image_path,
426
+ # rmbg_net=rmbg_net,
427
+ # seed=42,
428
+ # num_inference_steps=int(num_steps),
429
+ # guidance_scale=float(guidance_scale),
430
+ # faces=int(face_number),
431
+ # )
432
+
433
+ # if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
434
+ # raise ValueError("Mesh generation returned an empty mesh")
435
+
436
+ # # 🔧 Пересоздаём Trimesh и гарантируем чистоту геометрии
437
+ # mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, process=True)
438
+
439
+ # # ✅ Центрируе�� модель
440
+ # mesh.apply_translation(-mesh.center_mass)
441
+
442
+ # # ✅ Масштабируем к единичному размеру (все модели ~одинаковые)
443
+ # scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
444
+ # mesh.apply_scale(scale_factor)
445
+
446
+ # # ✅ Гарантированно пересчитываем нормали
447
+ # mesh.fix_normals()
448
+
449
+ # # print("[DEBUG] Normals present:", mesh.has_vertex_normals)
450
+ # if hasattr(mesh, "vertex_normals"):
451
+ # print("[DEBUG] Normals shape:", mesh.vertex_normals.shape)
452
+ # else:
453
+ # print("[DEBUG] Normals missing.")
454
+
455
+
456
+ # # 💾 Сохраняем GLB
457
+ # glb_data = mesh.export(file_type='glb')
458
+ # with open(output_path, "wb") as f:
459
+ # f.write(glb_data)
460
+
461
+ # print(f"[DEBUG] Mesh saved to {output_path}")
462
+ # return output_path if os.path.exists(output_path) else None
463
+
464
+ # except Exception as e:
465
+ # print("[ERROR]", e)
466
+ # traceback.print_exc()
467
+ # return f"Error: {e}"
468
+
469
+
470
+ # # Интерфейс Gradio
471
+ # demo = gr.Interface(
472
+ # fn=generate,
473
+ # inputs=gr.Image(type="filepath", label="Upload image"),
474
+ # outputs=gr.File(label="Download .glb"),
475
+ # title="TripoSG Image to 3D",
476
+ # description="Upload an image to generate a 3D model (.glb)",
477
+ # )
478
+
479
+ # # Запуск
480
+ # demo.launch()
481
+
482
+
483
+
484
+
485
+
486
+
487
+
488
+
489
+
490
+
491
  import os
492
  import subprocess
493
 
 
510
  import traceback
511
  import trimesh
512
  import numpy as np
513
+
514
+
515
  from trimesh.exchange.gltf import export_glb
516
 
517
  from inference_triposg import run_triposg
518
  from triposg.pipelines.pipeline_triposg import TripoSGPipeline
519
  from briarmbg import BriaRMBG
520
 
521
+ GLTF_PACK = "/tmp/gltfpack"
522
 
523
  print("Trimesh version:", trimesh.__version__)
524
 
 
554
  rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
555
  rmbg_net.eval()
556
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
557
 
558
  def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
559
  print("[API CALL] image_path received:", image_path)
 
596
  else:
597
  print("[DEBUG] Normals missing.")
598
 
 
599
  # 💾 Сохраняем GLB
600
  glb_data = mesh.export(file_type='glb')
601
  with open(output_path, "wb") as f:
 
619
  description="Upload an image to generate a 3D model (.glb)",
620
  )
621
 
622
+
623
+
624
  # Запуск
625
  demo.launch()
inference_triposg.py CHANGED
@@ -54,6 +54,7 @@ def run_triposg(
54
  num_inference_steps: int = 50,
55
  guidance_scale: float = 7.0,
56
  faces: int = -1,
 
57
  ) -> trimesh.Scene:
58
  print("[DEBUG] Preparing image...")
59
  img_pil = prepare_image(image_input, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
@@ -64,6 +65,7 @@ def run_triposg(
64
  generator=torch.Generator(device=pipe.device).manual_seed(seed),
65
  num_inference_steps=num_inference_steps,
66
  guidance_scale=guidance_scale,
 
67
  ).samples[0]
68
 
69
  print("[DEBUG] TripoSG output keys:", type(outputs), outputs[0].shape, outputs[1].shape)
 
54
  num_inference_steps: int = 50,
55
  guidance_scale: float = 7.0,
56
  faces: int = -1,
57
+ octree_depth: int = 9, # 👈 добавлено
58
  ) -> trimesh.Scene:
59
  print("[DEBUG] Preparing image...")
60
  img_pil = prepare_image(image_input, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
 
65
  generator=torch.Generator(device=pipe.device).manual_seed(seed),
66
  num_inference_steps=num_inference_steps,
67
  guidance_scale=guidance_scale,
68
+ flash_octree_depth=octree_depth, # 👈 добавлено
69
  ).samples[0]
70
 
71
  print("[DEBUG] TripoSG output keys:", type(outputs), outputs[0].shape, outputs[1].shape)