Spaces:
Sleeping
Sleeping
File size: 3,479 Bytes
b2a27a7 d23c9df f3bc318 f2228f5 b2a27a7 f3bc318 b2a27a7 b5ef988 b2a27a7 b5ef988 b2a27a7 d23c9df b2a27a7 d23c9df f3bc318 cf54aef d23c9df cf54aef f3bc318 f2228f5 b2a27a7 d23c9df f3bc318 b2a27a7 f3bc318 b2a27a7 |
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 |
import os
import subprocess
import gradio as gr
import uuid
import torch
import zipfile
import requests
import traceback
import trimesh
from trimesh.exchange.gltf import export_glb
from inference_triposg import run_triposg
from triposg.pipelines.pipeline_triposg import TripoSGPipeline
from briarmbg import BriaRMBG
# Убираем pyenv
os.environ.pop("PYENV_VERSION", None)
# Установка зависимостей
subprocess.run(["pip", "install", "torch", "wheel"], stdout=subprocess.DEVNULL)
subprocess.run([
"pip", "install", "--no-build-isolation",
"diso@git+https://github.com/SarahWeiii/diso.git"
], stdout=subprocess.DEVNULL)
print("Trimesh version:", trimesh.__version__)
# Настройки устройства
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
# Загрузка весов
weights_dir = "pretrained_weights"
triposg_path = os.path.join(weights_dir, "TripoSG")
rmbg_path = os.path.join(weights_dir, "RMBG-1.4")
if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
print("📦 Downloading pretrained weights...")
url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
zip_path = "pretrained_models.zip"
with requests.get(url, stream=True) as r:
r.raise_for_status()
with open(zip_path, "wb") as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
print("📦 Extracting weights...")
with zipfile.ZipFile(zip_path, "r") as zip_ref:
zip_ref.extractall(weights_dir)
os.remove(zip_path)
print("✅ Weights ready.")
# Загрузка моделей
pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
rmbg_net.eval()
# Генерация .glb
def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
print("[API CALL] image_path received:", image_path)
print("[API CALL] File exists:", os.path.exists(image_path))
temp_id = str(uuid.uuid4())
output_path = f"/tmp/{temp_id}.glb"
print("[DEBUG] Generating mesh from:", image_path)
try:
mesh = run_triposg(
pipe=pipe,
image_input=image_path,
rmbg_net=rmbg_net,
seed=42,
num_inference_steps=int(num_steps),
guidance_scale=float(guidance_scale),
faces=int(face_number),
)
if mesh is None:
raise ValueError("Mesh generation returned None")
# Очистка визуала, метаданных и имени
mesh.visual = None
mesh.metadata.clear()
mesh.name = "geometry_0"
# Экспорт в GLB без scene/world
glb_data = export_glb(mesh)
with open(output_path, "wb") as f:
f.write(glb_data)
print(f"[DEBUG] Mesh saved to {output_path}")
return output_path if os.path.exists(output_path) else None
except Exception as e:
print("[ERROR]", e)
traceback.print_exc()
return f"Error: {e}"
# Интерфейс Gradio
demo = gr.Interface(
fn=generate,
inputs=gr.Image(type="filepath", label="Upload image"),
outputs=gr.File(label="Download .glb"),
title="TripoSG Image to 3D",
description="Upload an image to generate a 3D model (.glb)",
)
# Запуск
demo.launch()
|