Spaces:
Running
on
Zero
Running
on
Zero
try that again last one was missing import spaces
Browse files
app.py
CHANGED
@@ -1,4 +1,163 @@
|
|
1 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
|
3 |
@spaces.GPU
|
4 |
def text_to_3d(
|
@@ -9,21 +168,27 @@ def text_to_3d(
|
|
9 |
slat_guidance_strength: float,
|
10 |
slat_sampling_steps: int,
|
11 |
req: gr.Request,
|
12 |
-
) -> Tuple[dict, str]: # Return type changed for
|
13 |
"""
|
14 |
Generates a 3D model (Gaussian and Mesh) from text and returns a
|
15 |
-
serializable state dictionary and a video preview path.
|
|
|
16 |
"""
|
17 |
-
print(f"[text_to_3d] Received prompt: '{prompt}', Seed: {seed}")
|
18 |
-
|
19 |
-
|
20 |
-
|
|
|
|
|
|
|
|
|
21 |
|
22 |
# --- Generation Pipeline ---
|
23 |
try:
|
24 |
-
print("[text_to_3d] Running Trellis pipeline...")
|
|
|
25 |
outputs = pipeline.run(
|
26 |
-
prompt,
|
27 |
seed=seed,
|
28 |
formats=["gaussian", "mesh"], # Ensure both are generated
|
29 |
sparse_structure_sampler_params={
|
@@ -34,54 +199,338 @@ def text_to_3d(
|
|
34 |
"steps": int(slat_sampling_steps), # Ensure steps are int
|
35 |
"cfg_strength": float(slat_guidance_strength),
|
36 |
},
|
|
|
37 |
)
|
38 |
-
print("[text_to_3d] Pipeline run completed.")
|
39 |
except Exception as e:
|
40 |
-
print(f"❌ [text_to_3d] Pipeline error: {e}", file=sys.stderr)
|
41 |
traceback.print_exc()
|
42 |
-
raise gr.Error(f"Trellis pipeline failed: {e}")
|
43 |
|
44 |
# --- Create Serializable State Dictionary ---
|
45 |
try:
|
46 |
state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
|
47 |
except Exception as e:
|
48 |
-
print(f"❌ [text_to_3d] pack_state error: {e}", file=sys.stderr)
|
49 |
traceback.print_exc()
|
50 |
-
raise gr.Error(f"Failed to pack state: {e}")
|
51 |
|
52 |
# --- Render Video Preview (TEMPORARILY DISABLED FOR DEBUGGING) ---
|
53 |
-
video_path = None #
|
|
|
|
|
54 |
# try:
|
55 |
# print("[text_to_3d] Rendering video preview...")
|
56 |
# video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
|
57 |
# video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
|
58 |
# # Ensure video frames are uint8
|
59 |
# video = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
|
60 |
-
#
|
61 |
-
# imageio.mimsave(
|
62 |
-
# print(f"[text_to_3d] Video saved to: {
|
|
|
63 |
# except Exception as e:
|
64 |
# print(f"❌ [text_to_3d] Video rendering/saving error: {e}", file=sys.stderr)
|
65 |
# traceback.print_exc()
|
66 |
# # Still return state_dict, but maybe signal video error? Return None for path.
|
67 |
# video_path = None # Indicate video failure
|
68 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
69 |
|
70 |
# --- Cleanup and Return ---
|
71 |
if torch.cuda.is_available():
|
72 |
torch.cuda.empty_cache()
|
73 |
-
print("[
|
|
|
|
|
|
|
|
|
|
|
74 |
|
75 |
-
# --- Return Serializable Dictionary and None for Video Path ---
|
76 |
-
print("[text_to_3d] Returning state dictionary and None video path.")
|
77 |
-
return state_dict, video_path # Return dict and None video path
|
78 |
|
79 |
# --- Gradio UI Definition ---
|
80 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
81 |
|
82 |
-
#
|
83 |
-
|
84 |
-
|
85 |
|
86 |
-
#
|
87 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Version: 1.1.2 - API State Fix + DEBUG (Video Disabled) + Import Fix (2025-05-04)
|
2 |
+
# Changes:
|
3 |
+
# - ENSURED `import spaces` is present for the @spaces.GPU decorator.
|
4 |
+
# - TEMPORARY DEBUGGING STEP: Commented out video rendering in `text_to_3d`
|
5 |
+
# and return None for video_path to isolate the "Session not found" error.
|
6 |
+
# - Modified `text_to_3d` to explicitly return the serializable `state_dict` from `pack_state`
|
7 |
+
# as the first return value. This ensures the dictionary is available via the API.
|
8 |
+
# - Modified `extract_glb` and `extract_gaussian` to accept `state_dict: dict` as their first argument
|
9 |
+
# instead of relying on the implicit `gr.State` object type when called via API.
|
10 |
+
# - Kept Gradio UI bindings (`outputs=[output_buf, ...]`, `inputs=[output_buf, ...]`)
|
11 |
+
# so the UI continues to function by passing the dictionary through output_buf.
|
12 |
+
# - Added minor safety checks and logging.
|
13 |
+
|
14 |
+
import gradio as gr
|
15 |
+
import spaces # <<<--- ENSURE THIS IMPORT IS PRESENT
|
16 |
+
|
17 |
+
import os
|
18 |
+
import shutil
|
19 |
+
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
|
20 |
+
# Fix potential SpConv issue if needed, try 'hash' or 'native'
|
21 |
+
# os.environ.setdefault('SPCONV_ALGO', 'native') # Use setdefault to avoid overwriting if already set
|
22 |
+
os.environ['SPCONV_ALGO'] = 'native' # Direct set as per original
|
23 |
+
|
24 |
+
from typing import *
|
25 |
+
import torch
|
26 |
+
import numpy as np
|
27 |
+
import imageio
|
28 |
+
from easydict import EasyDict as edict
|
29 |
+
from trellis.pipelines import TrellisTextTo3DPipeline
|
30 |
+
from trellis.representations import Gaussian, MeshExtractResult
|
31 |
+
from trellis.utils import render_utils, postprocessing_utils
|
32 |
+
|
33 |
+
import traceback
|
34 |
+
import sys
|
35 |
+
|
36 |
+
|
37 |
+
MAX_SEED = np.iinfo(np.int32).max
|
38 |
+
# Ensure TMP_DIR is correctly defined relative to the script location
|
39 |
+
# Using /tmp/ directly might be more robust in some container environments
|
40 |
+
# TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
|
41 |
+
TMP_DIR = '/tmp/gradio_sessions' # Use standard /tmp directory
|
42 |
+
print(f"Using temporary directory: {TMP_DIR}")
|
43 |
+
os.makedirs(TMP_DIR, exist_ok=True)
|
44 |
+
|
45 |
+
|
46 |
+
def start_session(req: gr.Request):
|
47 |
+
"""Creates a temporary directory for the user session."""
|
48 |
+
try:
|
49 |
+
session_hash = req.session_hash
|
50 |
+
if not session_hash:
|
51 |
+
# Fallback or generate a temporary ID if session_hash is missing (might happen on first load?)
|
52 |
+
session_hash = f"no_session_{np.random.randint(10000, 99999)}"
|
53 |
+
print(f"Warning: No session_hash in request, using temporary ID: {session_hash}")
|
54 |
+
|
55 |
+
user_dir = os.path.join(TMP_DIR, str(session_hash))
|
56 |
+
os.makedirs(user_dir, exist_ok=True)
|
57 |
+
print(f"Started session, created directory: {user_dir}")
|
58 |
+
except Exception as e:
|
59 |
+
print(f"Error in start_session: {e}", file=sys.stderr)
|
60 |
+
# Decide if this is critical - maybe raise to prevent further issues?
|
61 |
+
|
62 |
+
|
63 |
+
def end_session(req: gr.Request):
|
64 |
+
"""Removes the temporary directory for the user session."""
|
65 |
+
try:
|
66 |
+
session_hash = req.session_hash
|
67 |
+
if not session_hash:
|
68 |
+
print("Warning: No session_hash in end_session request, cannot clean up.")
|
69 |
+
return
|
70 |
+
|
71 |
+
user_dir = os.path.join(TMP_DIR, str(session_hash))
|
72 |
+
if os.path.exists(user_dir):
|
73 |
+
try:
|
74 |
+
shutil.rmtree(user_dir)
|
75 |
+
print(f"Ended session, removed directory: {user_dir}")
|
76 |
+
except OSError as e:
|
77 |
+
print(f"Error removing tmp directory {user_dir}: {e.strerror}", file=sys.stderr)
|
78 |
+
else:
|
79 |
+
print(f"Ended session, directory already removed or hash mismatch: {user_dir}")
|
80 |
+
except Exception as e:
|
81 |
+
print(f"Error in end_session: {e}", file=sys.stderr)
|
82 |
+
|
83 |
+
|
84 |
+
def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
|
85 |
+
"""Packs Gaussian and Mesh data into a serializable dictionary."""
|
86 |
+
print("[pack_state] Packing state to dictionary...")
|
87 |
+
try:
|
88 |
+
packed_data = {
|
89 |
+
'gaussian': {
|
90 |
+
**{k: v for k, v in gs.init_params.items()}, # Ensure init_params are included
|
91 |
+
'_xyz': gs._xyz.detach().cpu().numpy(),
|
92 |
+
'_features_dc': gs._features_dc.detach().cpu().numpy(),
|
93 |
+
'_scaling': gs._scaling.detach().cpu().numpy(),
|
94 |
+
'_rotation': gs._rotation.detach().cpu().numpy(),
|
95 |
+
'_opacity': gs._opacity.detach().cpu().numpy(),
|
96 |
+
},
|
97 |
+
'mesh': {
|
98 |
+
'vertices': mesh.vertices.detach().cpu().numpy(),
|
99 |
+
'faces': mesh.faces.detach().cpu().numpy(),
|
100 |
+
},
|
101 |
+
}
|
102 |
+
print(f"[pack_state] Dictionary created. Keys: {list(packed_data.keys())}, Gaussian points: {len(packed_data['gaussian']['_xyz'])}, Mesh vertices: {len(packed_data['mesh']['vertices'])}")
|
103 |
+
return packed_data
|
104 |
+
except Exception as e:
|
105 |
+
print(f"Error during pack_state: {e}", file=sys.stderr)
|
106 |
+
traceback.print_exc()
|
107 |
+
raise # Re-raise the error to be caught upstream
|
108 |
+
|
109 |
+
|
110 |
+
def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]:
|
111 |
+
"""Unpacks Gaussian and Mesh data from a dictionary."""
|
112 |
+
print("[unpack_state] Unpacking state from dictionary...")
|
113 |
+
try:
|
114 |
+
if not isinstance(state_dict, dict) or 'gaussian' not in state_dict or 'mesh' not in state_dict:
|
115 |
+
raise ValueError("Invalid state_dict structure passed to unpack_state.")
|
116 |
+
|
117 |
+
# Ensure the device is correctly set when unpacking
|
118 |
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
119 |
+
print(f"[unpack_state] Using device: {device}")
|
120 |
+
|
121 |
+
gauss_data = state_dict['gaussian']
|
122 |
+
mesh_data = state_dict['mesh']
|
123 |
+
|
124 |
+
# Recreate Gaussian object using parameters stored during packing
|
125 |
+
gs = Gaussian(
|
126 |
+
aabb=gauss_data.get('aabb'), # Use .get for safety
|
127 |
+
sh_degree=gauss_data.get('sh_degree'),
|
128 |
+
mininum_kernel_size=gauss_data.get('mininum_kernel_size'),
|
129 |
+
scaling_bias=gauss_data.get('scaling_bias'),
|
130 |
+
opacity_bias=gauss_data.get('opacity_bias'),
|
131 |
+
scaling_activation=gauss_data.get('scaling_activation'),
|
132 |
+
)
|
133 |
+
# Load tensors, ensuring they are created on the correct device
|
134 |
+
gs._xyz = torch.tensor(gauss_data['_xyz'], device=device, dtype=torch.float32)
|
135 |
+
gs._features_dc = torch.tensor(gauss_data['_features_dc'], device=device, dtype=torch.float32)
|
136 |
+
gs._scaling = torch.tensor(gauss_data['_scaling'], device=device, dtype=torch.float32)
|
137 |
+
gs._rotation = torch.tensor(gauss_data['_rotation'], device=device, dtype=torch.float32)
|
138 |
+
gs._opacity = torch.tensor(gauss_data['_opacity'], device=device, dtype=torch.float32)
|
139 |
+
print(f"[unpack_state] Gaussian unpacked. Points: {gs.get_xyz.shape[0]}")
|
140 |
+
|
141 |
+
# Recreate mesh object using edict for compatibility if needed elsewhere
|
142 |
+
mesh = edict(
|
143 |
+
vertices=torch.tensor(mesh_data['vertices'], device=device, dtype=torch.float32),
|
144 |
+
faces=torch.tensor(mesh_data['faces'], device=device, dtype=torch.int64), # Faces are typically long/int64
|
145 |
+
)
|
146 |
+
print(f"[unpack_state] Mesh unpacked. Vertices: {mesh.vertices.shape[0]}, Faces: {mesh.faces.shape[0]}")
|
147 |
+
|
148 |
+
return gs, mesh
|
149 |
+
except Exception as e:
|
150 |
+
print(f"Error during unpack_state: {e}", file=sys.stderr)
|
151 |
+
traceback.print_exc()
|
152 |
+
raise # Re-raise the error
|
153 |
+
|
154 |
+
|
155 |
+
def get_seed(randomize_seed: bool, seed: int) -> int:
|
156 |
+
"""Gets a seed value, randomizing if requested."""
|
157 |
+
new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
|
158 |
+
print(f"[get_seed] Randomize: {randomize_seed}, Input Seed: {seed}, Output Seed: {new_seed}")
|
159 |
+
return int(new_seed) # Ensure it's a standard int
|
160 |
+
|
161 |
|
162 |
@spaces.GPU
|
163 |
def text_to_3d(
|
|
|
168 |
slat_guidance_strength: float,
|
169 |
slat_sampling_steps: int,
|
170 |
req: gr.Request,
|
171 |
+
) -> Tuple[dict, Optional[str]]: # Return type changed Optional[str] for video path
|
172 |
"""
|
173 |
Generates a 3D model (Gaussian and Mesh) from text and returns a
|
174 |
+
serializable state dictionary and potentially a video preview path.
|
175 |
+
>>> TEMPORARILY DISABLED VIDEO RENDERING FOR DEBUGGING <<<
|
176 |
"""
|
177 |
+
print(f"[text_to_3d - DEBUG MODE] Received prompt: '{prompt}', Seed: {seed}")
|
178 |
+
session_hash = req.session_hash
|
179 |
+
if not session_hash:
|
180 |
+
session_hash = f"no_session_{np.random.randint(10000, 99999)}" # Use consistent fallback
|
181 |
+
print(f"Warning: No session_hash in text_to_3d request, using temporary ID: {session_hash}")
|
182 |
+
user_dir = os.path.join(TMP_DIR, str(session_hash))
|
183 |
+
os.makedirs(user_dir, exist_ok=True) # Ensure it exists for this request
|
184 |
+
print(f"[text_to_3d - DEBUG MODE] User directory: {user_dir}")
|
185 |
|
186 |
# --- Generation Pipeline ---
|
187 |
try:
|
188 |
+
print("[text_to_3d - DEBUG MODE] Running Trellis pipeline...")
|
189 |
+
# Add more specific pipeline settings if needed based on Trellis docs
|
190 |
outputs = pipeline.run(
|
191 |
+
prompt=prompt,
|
192 |
seed=seed,
|
193 |
formats=["gaussian", "mesh"], # Ensure both are generated
|
194 |
sparse_structure_sampler_params={
|
|
|
199 |
"steps": int(slat_sampling_steps), # Ensure steps are int
|
200 |
"cfg_strength": float(slat_guidance_strength),
|
201 |
},
|
202 |
+
# device='cuda' # Explicitly specify device if needed
|
203 |
)
|
204 |
+
print("[text_to_3d - DEBUG MODE] Pipeline run completed.")
|
205 |
except Exception as e:
|
206 |
+
print(f"❌ [text_to_3d - DEBUG MODE] Pipeline error: {e}", file=sys.stderr)
|
207 |
traceback.print_exc()
|
208 |
+
raise gr.Error(f"Trellis pipeline failed during generation: {e}") # More specific error
|
209 |
|
210 |
# --- Create Serializable State Dictionary ---
|
211 |
try:
|
212 |
state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
|
213 |
except Exception as e:
|
214 |
+
print(f"❌ [text_to_3d - DEBUG MODE] pack_state error: {e}", file=sys.stderr)
|
215 |
traceback.print_exc()
|
216 |
+
raise gr.Error(f"Failed to pack state after generation: {e}")
|
217 |
|
218 |
# --- Render Video Preview (TEMPORARILY DISABLED FOR DEBUGGING) ---
|
219 |
+
video_path = None # Explicitly set path to None for this debug version
|
220 |
+
print("[text_to_3d - DEBUG MODE] Skipping video rendering.")
|
221 |
+
# --- Original Video Code Block Start (Keep commented for now) ---
|
222 |
# try:
|
223 |
# print("[text_to_3d] Rendering video preview...")
|
224 |
# video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
|
225 |
# video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
|
226 |
# # Ensure video frames are uint8
|
227 |
# video = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
|
228 |
+
# video_path_tmp = os.path.join(user_dir, 'sample.mp4') # Use temp name
|
229 |
+
# imageio.mimsave(video_path_tmp, video, fps=15, quality=8) # Added quality setting
|
230 |
+
# print(f"[text_to_3d] Video saved to: {video_path_tmp}")
|
231 |
+
# video_path = video_path_tmp # Assign if successful
|
232 |
# except Exception as e:
|
233 |
# print(f"❌ [text_to_3d] Video rendering/saving error: {e}", file=sys.stderr)
|
234 |
# traceback.print_exc()
|
235 |
# # Still return state_dict, but maybe signal video error? Return None for path.
|
236 |
# video_path = None # Indicate video failure
|
237 |
+
# --- Original Video Code Block End ---
|
238 |
+
|
239 |
+
# --- Cleanup and Return ---
|
240 |
+
if torch.cuda.is_available():
|
241 |
+
torch.cuda.empty_cache()
|
242 |
+
print("[text_to_3d - DEBUG MODE] Cleared CUDA cache.")
|
243 |
+
|
244 |
+
# --- Return Serializable Dictionary and None Video Path ---
|
245 |
+
print("[text_to_3d - DEBUG MODE] Returning state dictionary and None video path.")
|
246 |
+
# Ensure state_dict is not None before returning
|
247 |
+
if state_dict is None:
|
248 |
+
raise gr.Error("Failed to create state dictionary.")
|
249 |
+
return state_dict, video_path
|
250 |
+
|
251 |
+
|
252 |
+
@spaces.GPU(duration=120) # Increased duration slightly
|
253 |
+
def extract_glb(
|
254 |
+
state_dict: dict, # <-- Accepts the dictionary directly
|
255 |
+
mesh_simplify: float,
|
256 |
+
texture_size: int,
|
257 |
+
req: gr.Request,
|
258 |
+
) -> Tuple[str, str]:
|
259 |
+
"""
|
260 |
+
Extracts a GLB file from the provided 3D model state dictionary.
|
261 |
+
"""
|
262 |
+
print(f"[extract_glb] Received request. Simplify: {mesh_simplify}, Texture Size: {texture_size}")
|
263 |
+
session_hash = req.session_hash
|
264 |
+
if not session_hash:
|
265 |
+
session_hash = f"no_session_{np.random.randint(10000, 99999)}"
|
266 |
+
print(f"Warning: No session_hash in extract_glb request, using temporary ID: {session_hash}")
|
267 |
+
|
268 |
+
if not isinstance(state_dict, dict):
|
269 |
+
print("❌ [extract_glb] Error: Invalid state_dict received (not a dictionary).")
|
270 |
+
raise gr.Error("Invalid state data received. Please generate the model first.")
|
271 |
+
|
272 |
+
user_dir = os.path.join(TMP_DIR, str(session_hash))
|
273 |
+
os.makedirs(user_dir, exist_ok=True) # Ensure it exists
|
274 |
+
print(f"[extract_glb] User directory: {user_dir}")
|
275 |
+
|
276 |
+
# --- Unpack state from the dictionary ---
|
277 |
+
try:
|
278 |
+
gs, mesh = unpack_state(state_dict)
|
279 |
+
except Exception as e:
|
280 |
+
print(f"❌ [extract_glb] unpack_state error: {e}", file=sys.stderr)
|
281 |
+
traceback.print_exc()
|
282 |
+
raise gr.Error(f"Failed to unpack state during GLB extraction: {e}")
|
283 |
+
|
284 |
+
# --- Postprocessing and Export ---
|
285 |
+
try:
|
286 |
+
print("[extract_glb] Converting to GLB...")
|
287 |
+
# Ensure parameters have correct types
|
288 |
+
simplify_factor = float(mesh_simplify)
|
289 |
+
tex_size = int(texture_size)
|
290 |
+
glb = postprocessing_utils.to_glb(gs, mesh, simplify=simplify_factor, texture_size=tex_size, verbose=True)
|
291 |
+
glb_path = os.path.join(user_dir, 'sample.glb')
|
292 |
+
print(f"[extract_glb] Exporting GLB to: {glb_path}")
|
293 |
+
glb.export(glb_path)
|
294 |
+
print("[extract_glb] GLB exported successfully.")
|
295 |
+
except Exception as e:
|
296 |
+
print(f"❌ [extract_glb] GLB conversion/export error: {e}", file=sys.stderr)
|
297 |
+
traceback.print_exc()
|
298 |
+
raise gr.Error(f"Failed to extract GLB: {e}")
|
299 |
+
|
300 |
+
# --- Cleanup and Return ---
|
301 |
+
if torch.cuda.is_available():
|
302 |
+
torch.cuda.empty_cache()
|
303 |
+
print("[extract_glb] Cleared CUDA cache.")
|
304 |
+
|
305 |
+
# Return path twice for both Model3D and DownloadButton components
|
306 |
+
print("[extract_glb] Returning GLB path.")
|
307 |
+
# Ensure path is returned, even if export failed somehow (though error should raise first)
|
308 |
+
return glb_path, glb_path
|
309 |
+
|
310 |
+
|
311 |
+
@spaces.GPU
|
312 |
+
def extract_gaussian(
|
313 |
+
state_dict: dict, # <-- Accepts the dictionary directly
|
314 |
+
req: gr.Request
|
315 |
+
) -> Tuple[str, str]:
|
316 |
+
"""
|
317 |
+
Extracts a PLY (Gaussian) file from the provided 3D model state dictionary.
|
318 |
+
"""
|
319 |
+
print("[extract_gaussian] Received request.")
|
320 |
+
session_hash = req.session_hash
|
321 |
+
if not session_hash:
|
322 |
+
session_hash = f"no_session_{np.random.randint(10000, 99999)}"
|
323 |
+
print(f"Warning: No session_hash in extract_gaussian request, using temporary ID: {session_hash}")
|
324 |
+
|
325 |
+
if not isinstance(state_dict, dict):
|
326 |
+
print("❌ [extract_gaussian] Error: Invalid state_dict received (not a dictionary).")
|
327 |
+
raise gr.Error("Invalid state data received. Please generate the model first.")
|
328 |
+
|
329 |
+
user_dir = os.path.join(TMP_DIR, str(session_hash))
|
330 |
+
os.makedirs(user_dir, exist_ok=True) # Ensure it exists
|
331 |
+
print(f"[extract_gaussian] User directory: {user_dir}")
|
332 |
+
|
333 |
+
# --- Unpack state from the dictionary ---
|
334 |
+
try:
|
335 |
+
gs, _ = unpack_state(state_dict) # Only need Gaussian part
|
336 |
+
except Exception as e:
|
337 |
+
print(f"❌ [extract_gaussian] unpack_state error: {e}", file=sys.stderr)
|
338 |
+
traceback.print_exc()
|
339 |
+
raise gr.Error(f"Failed to unpack state during Gaussian extraction: {e}")
|
340 |
+
|
341 |
+
# --- Export PLY ---
|
342 |
+
try:
|
343 |
+
gaussian_path = os.path.join(user_dir, 'sample.ply')
|
344 |
+
print(f"[extract_gaussian] Saving PLY to: {gaussian_path}")
|
345 |
+
gs.save_ply(gaussian_path)
|
346 |
+
print("[extract_gaussian] PLY saved successfully.")
|
347 |
+
except Exception as e:
|
348 |
+
print(f"❌ [extract_gaussian] PLY saving error: {e}", file=sys.stderr)
|
349 |
+
traceback.print_exc()
|
350 |
+
raise gr.Error(f"Failed to extract Gaussian PLY: {e}")
|
351 |
|
352 |
# --- Cleanup and Return ---
|
353 |
if torch.cuda.is_available():
|
354 |
torch.cuda.empty_cache()
|
355 |
+
print("[extract_gaussian] Cleared CUDA cache.")
|
356 |
+
|
357 |
+
# Return path twice for both Model3D and DownloadButton components
|
358 |
+
print("[extract_gaussian] Returning PLY path.")
|
359 |
+
# Ensure path is returned
|
360 |
+
return gaussian_path, gaussian_path
|
361 |
|
|
|
|
|
|
|
362 |
|
363 |
# --- Gradio UI Definition ---
|
364 |
+
print("Setting up Gradio Blocks interface...")
|
365 |
+
# Define the interface layout
|
366 |
+
with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
|
367 |
+
gr.Markdown("""
|
368 |
+
# Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
|
369 |
+
* Type a text prompt and click "Generate" to create a 3D asset preview.
|
370 |
+
* Adjust extraction settings if desired.
|
371 |
+
* Click "Extract GLB" or "Extract Gaussian" to get the downloadable 3D file.
|
372 |
+
*(Note: Video preview is temporarily disabled for debugging)*
|
373 |
+
""")
|
374 |
+
|
375 |
+
# --- State Buffer ---
|
376 |
+
# This hidden component holds the dictionary linking generation and extraction.
|
377 |
+
output_buf = gr.State()
|
378 |
+
|
379 |
+
with gr.Row():
|
380 |
+
with gr.Column(scale=1): # Input column
|
381 |
+
text_prompt = gr.Textbox(label="Text Prompt", lines=5, placeholder="e.g., a cute red dragon")
|
382 |
+
|
383 |
+
with gr.Accordion(label="Generation Settings", open=False):
|
384 |
+
seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
|
385 |
+
randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
|
386 |
+
gr.Markdown("--- \n **Stage 1: Sparse Structure Generation**")
|
387 |
+
with gr.Row():
|
388 |
+
ss_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
|
389 |
+
ss_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
|
390 |
+
gr.Markdown("--- \n **Stage 2: Structured Latent Generation**")
|
391 |
+
with gr.Row():
|
392 |
+
slat_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
|
393 |
+
slat_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
|
394 |
+
|
395 |
+
generate_btn = gr.Button("Generate 3D Preview", variant="primary")
|
396 |
+
|
397 |
+
with gr.Accordion(label="GLB Extraction Settings", open=True): # Open by default
|
398 |
+
mesh_simplify = gr.Slider(0.9, 0.99, label="Simplify Factor", value=0.95, step=0.01, info="Higher value = less simplification (more polys)")
|
399 |
+
texture_size = gr.Slider(512, 2048, label="Texture Size (pixels)", value=1024, step=512, info="Size of the generated texture map")
|
400 |
+
|
401 |
+
with gr.Row():
|
402 |
+
extract_glb_btn = gr.Button("Extract GLB", interactive=False)
|
403 |
+
extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False)
|
404 |
+
gr.Markdown("""
|
405 |
+
*NOTE: Gaussian file (.ply) can be very large (~50MB+) and may take time to process/download.*
|
406 |
+
""")
|
407 |
+
|
408 |
+
with gr.Column(scale=1): # Output column
|
409 |
+
# Video component remains for layout but won't show anything in this debug version
|
410 |
+
video_output = gr.Video(label="Generated 3D Preview (DISABLED FOR DEBUG)", autoplay=False, loop=False, value=None, height=350)
|
411 |
+
model_output = gr.Model3D(label="Extracted Model Preview", height=350, clear_color=[0.95, 0.95, 0.95, 1.0]) # Light background
|
412 |
+
|
413 |
+
with gr.Row():
|
414 |
+
download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
|
415 |
+
download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False)
|
416 |
+
|
417 |
+
# --- Event Handlers ---
|
418 |
+
print("Defining Gradio event handlers...")
|
419 |
|
420 |
+
# Handle session start/end
|
421 |
+
demo.load(start_session, inputs=None, outputs=None)
|
422 |
+
demo.unload(end_session, inputs=None, outputs=None)
|
423 |
|
424 |
+
# --- Generate Button Click Flow ---
|
425 |
+
generate_event = generate_btn.click(
|
426 |
+
get_seed,
|
427 |
+
inputs=[randomize_seed, seed],
|
428 |
+
outputs=[seed],
|
429 |
+
api_name="get_seed"
|
430 |
+
).then(
|
431 |
+
text_to_3d,
|
432 |
+
inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
|
433 |
+
# Output state_dict to buffer, output None to video component
|
434 |
+
outputs=[output_buf, video_output],
|
435 |
+
api_name="text_to_3d"
|
436 |
+
).then(
|
437 |
+
# Function to update button interactivity after generation attempt
|
438 |
+
lambda: (
|
439 |
+
gr.Button(interactive=True),
|
440 |
+
gr.Button(interactive=True),
|
441 |
+
gr.DownloadButton(interactive=False),
|
442 |
+
gr.DownloadButton(interactive=False)
|
443 |
+
),
|
444 |
+
inputs=None, # No inputs needed for the lambda
|
445 |
+
outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
|
446 |
+
)
|
447 |
+
|
448 |
+
# --- Extract GLB Button Click Flow ---
|
449 |
+
extract_glb_event = extract_glb_btn.click(
|
450 |
+
extract_glb,
|
451 |
+
inputs=[output_buf, mesh_simplify, texture_size],
|
452 |
+
outputs=[model_output, download_glb],
|
453 |
+
api_name="extract_glb"
|
454 |
+
).then(
|
455 |
+
lambda: gr.DownloadButton(interactive=True),
|
456 |
+
inputs=None,
|
457 |
+
outputs=[download_glb],
|
458 |
+
)
|
459 |
+
|
460 |
+
# --- Extract Gaussian Button Click Flow ---
|
461 |
+
extract_gs_event = extract_gs_btn.click(
|
462 |
+
extract_gaussian,
|
463 |
+
inputs=[output_buf],
|
464 |
+
outputs=[model_output, download_gs],
|
465 |
+
api_name="extract_gaussian"
|
466 |
+
).then(
|
467 |
+
lambda: gr.DownloadButton(interactive=True),
|
468 |
+
inputs=None,
|
469 |
+
outputs=[download_gs],
|
470 |
+
)
|
471 |
+
|
472 |
+
# --- Clear Download Button Interactivity when model preview is cleared ---
|
473 |
+
model_output.clear(
|
474 |
+
lambda: (gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)),
|
475 |
+
inputs=None,
|
476 |
+
outputs=[download_glb, download_gs]
|
477 |
+
)
|
478 |
+
# Also disable buttons if the (currently disabled) video output is cleared
|
479 |
+
video_output.clear(
|
480 |
+
lambda: (
|
481 |
+
gr.Button(interactive=False),
|
482 |
+
gr.Button(interactive=False),
|
483 |
+
gr.DownloadButton(interactive=False),
|
484 |
+
gr.DownloadButton(interactive=False)
|
485 |
+
),
|
486 |
+
inputs=None,
|
487 |
+
outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
|
488 |
+
)
|
489 |
+
|
490 |
+
print("Gradio interface setup complete.")
|
491 |
+
|
492 |
+
|
493 |
+
# --- Launch the Gradio app ---
|
494 |
+
# Main execution block
|
495 |
+
if __name__ == "__main__":
|
496 |
+
print("Loading Trellis pipeline...")
|
497 |
+
pipeline_loaded = False
|
498 |
+
try:
|
499 |
+
# Ensure model/variant matches requirements, use revision if needed
|
500 |
+
pipeline = TrellisTextTo3DPipeline.from_pretrained(
|
501 |
+
"JeffreyXiang/TRELLIS-text-xlarge",
|
502 |
+
# revision="main", # Specify if needed
|
503 |
+
torch_dtype=torch.float16 # Use float16 if GPU supports it for less memory
|
504 |
+
)
|
505 |
+
# Move to GPU if available
|
506 |
+
if torch.cuda.is_available():
|
507 |
+
pipeline = pipeline.to("cuda")
|
508 |
+
print("✅ Trellis pipeline loaded successfully to GPU.")
|
509 |
+
else:
|
510 |
+
print("⚠️ WARNING: CUDA not available, running on CPU (will be very slow).")
|
511 |
+
print("✅ Trellis pipeline loaded successfully to CPU.")
|
512 |
+
pipeline_loaded = True
|
513 |
+
except Exception as e:
|
514 |
+
print(f"❌ Failed to load Trellis pipeline: {e}", file=sys.stderr)
|
515 |
+
traceback.print_exc()
|
516 |
+
# Exit if pipeline is critical for the app to run
|
517 |
+
print("❌ Exiting due to pipeline load failure.")
|
518 |
+
sys.exit(1) # Exit if pipeline fails
|
519 |
+
|
520 |
+
if pipeline_loaded:
|
521 |
+
print("Launching Gradio demo...")
|
522 |
+
# Set share=True if you need a public link (e.g., for testing from outside local network)
|
523 |
+
# Set server_name="0.0.0.0" to allow access from local network IP
|
524 |
+
# Increased concurrency_limit and timeout for queue might help
|
525 |
+
demo.queue(
|
526 |
+
# default_concurrency_limit=5, # Adjust based on expected load and space resources
|
527 |
+
# api_open=True # Keep API accessible
|
528 |
+
).launch(
|
529 |
+
# server_name="0.0.0.0", # Make accessible on local network
|
530 |
+
# share=False, # Set to True for public link if needed
|
531 |
+
debug=True, # Enable Gradio debug mode for more detailed logs
|
532 |
+
# prevent_thread_lock=True # May help with async issues in some cases
|
533 |
+
)
|
534 |
+
print("Gradio demo launched.")
|
535 |
+
else:
|
536 |
+
print("Gradio demo not launched due to pipeline loading failure.")
|