Update app.py

app.py

@@ -1,11 +1,11 @@
-# Version: 1.1.2 - Removed torch_dtype from from_pretrained call
-# Applied:
+# Version: 1.1.3 - Load pipeline at module level for Spaces environment
+# Applied targeted fixes:
 #  - Removed unsupported inputs/outputs kwargs on demo.load/unload
 #  - Converted NumPy arrays to lists in pack_state for JSON safety
 #  - Fixed indentation in Blocks event-handlers
 #  - Verified clear() callbacks use only callback + outputs
-#  - Removed `torch_dtype` arg from TrellisTextTo3DPipeline.from_pretrained
-#  - Bumped version, added comments at change sites
+#  - Removed `torch_dtype` arg from from_pretrained
+#  - Moved pipeline initialization to module level so it's available in threads
 
 import gradio as gr
 import spaces
@@ -25,10 +25,26 @@ from trellis.utils import render_utils, postprocessing_utils
 import traceback
 import sys
 
+# --- Global Config ---
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
+# --- Initialize Trellis Pipeline at import time ---
+print("[Startup] Loading Trellis pipeline...")
+try:
+    pipeline = TrellisTextTo3DPipeline.from_pretrained(
+        "JeffreyXiang/TRELLIS-text-xlarge"
+    )
+    if torch.cuda.is_available():
+        pipeline = pipeline.to("cuda")
+        print("[Startup] Trellis pipeline loaded to GPU.")
+    else:
+        print("[Startup] Trellis pipeline loaded to CPU.")
+except Exception as e:
+    print(f"❌ [Startup] Failed to load Trellis pipeline: {e}")
+    raise
+
 
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -50,10 +66,9 @@ def end_session(req: gr.Request):
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     """Packs Gaussian and Mesh data into a JSON-serializable dictionary."""
-    packed_data = {
+    return {
         'gaussian': {
             **{k: v for k, v in gs.init_params.items()},
-            # FIX: convert arrays to lists for JSON
             '_xyz': gs._xyz.detach().cpu().numpy().tolist(),
             '_features_dc': gs._features_dc.detach().cpu().numpy().tolist(),
             '_scaling': gs._scaling.detach().cpu().numpy().tolist(),
@@ -65,178 +80,128 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
             'faces': mesh.faces.detach().cpu().numpy().tolist(),
         },
     }
-    return packed_data
 
 
 def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]:
-    print("[unpack_state] Unpacking state from dictionary... ")
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    gauss_data = state_dict['gaussian']
-    mesh_data = state_dict['mesh']
+    gd = state_dict['gaussian']
+    md = state_dict['mesh']
     gs = Gaussian(
-        aabb=gauss_data.get('aabb'),
-        sh_degree=gauss_data.get('sh_degree'),
-        mininum_kernel_size=gauss_data.get('mininum_kernel_size'),
-        scaling_bias=gauss_data.get('scaling_bias'),
-        opacity_bias=gauss_data.get('opacity_bias'),
-        scaling_activation=gauss_data.get('scaling_activation'),
+        aabb=gd.get('aabb'), sh_degree=gd.get('sh_degree'),
+        mininum_kernel_size=gd.get('mininum_kernel_size'),
+        scaling_bias=gd.get('scaling_bias'), opacity_bias=gd.get('opacity_bias'),
+        scaling_activation=gd.get('scaling_activation')
     )
-    gs._xyz = torch.tensor(np.array(gauss_data['_xyz']), device=device, dtype=torch.float32)
-    gs._features_dc = torch.tensor(np.array(gauss_data['_features_dc']), device=device, dtype=torch.float32)
-    gs._scaling = torch.tensor(np.array(gauss_data['_scaling']), device=device, dtype=torch.float32)
-    gs._rotation = torch.tensor(np.array(gauss_data['_rotation']), device=device, dtype=torch.float32)
-    gs._opacity = torch.tensor(np.array(gauss_data['_opacity']), device=device, dtype=torch.float32)
+    gs._xyz = torch.tensor(np.array(gd['_xyz']), device=device, dtype=torch.float32)
+    gs._features_dc = torch.tensor(np.array(gd['_features_dc']), device=device, dtype=torch.float32)
+    gs._scaling = torch.tensor(np.array(gd['_scaling']), device=device, dtype=torch.float32)
+    gs._rotation = torch.tensor(np.array(gd['_rotation']), device=device, dtype=torch.float32)
+    gs._opacity = torch.tensor(np.array(gd['_opacity']), device=device, dtype=torch.float32)
     mesh = edict(
-        vertices=torch.tensor(np.array(mesh_data['vertices']), device=device, dtype=torch.float32),
-        faces=torch.tensor(np.array(mesh_data['faces']), device=device, dtype=torch.int64),
+        vertices=torch.tensor(np.array(md['vertices']), device=device, dtype=torch.float32),
+        faces=torch.tensor(np.array(md['faces']), device=device, dtype=torch.int64),
     )
     return gs, mesh
 
 
 def get_seed(randomize_seed: bool, seed: int) -> int:
-    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
-    return int(new_seed)
+    return int(np.random.randint(0, MAX_SEED) if randomize_seed else seed)
 
 @spaces.GPU
 def text_to_3d(
-    prompt: str,
-    seed: int,
-    ss_guidance_strength: float,
-    ss_sampling_steps: int,
-    slat_guidance_strength: float,
-    slat_sampling_steps: int,
-    req: gr.Request,
+    prompt: str, seed: int,
+    ss_guidance_strength: float, ss_sampling_steps: int,
+    slat_guidance_strength: float, slat_sampling_steps: int,
+    req: gr.Request
 ) -> Tuple[dict, str]:
-    outputs = pipeline.run(
-        prompt,
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        sparse_structure_sampler_params={"steps": int(ss_sampling_steps), "cfg_strength": float(ss_guidance_strength)},
-        slat_sampler_params={"steps": int(slat_sampling_steps), "cfg_strength": float(slat_guidance_strength)},
+    out = pipeline.run(
+        prompt, seed=seed,
+        formats=["gaussian","mesh"],
+        sparse_structure_sampler_params={"steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength},
+        slat_sampler_params={"steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength}
     )
-    state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-    video_combined = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-    video_path = os.path.join(user_dir, 'sample.mp4')
-    imageio.mimsave(video_path, video_combined, fps=15, quality=8)
+    state = pack_state(out['gaussian'][0], out['mesh'][0])
+    vid_c = render_utils.render_video(out['gaussian'][0],num_frames=120)['color']
+    vid_n = render_utils.render_video(out['mesh'][0],num_frames=120)['normal']
+    vid = [np.concatenate([c.astype(np.uint8), n.astype(np.uint8)], axis=1) for c,n in zip(vid_c,vid_n)]
+    ud = os.path.join(TMP_DIR,str(req.session_hash)); os.makedirs(ud,exist_ok=True)
+    vp = os.path.join(ud,'sample.mp4'); imageio.mimsave(vp,vid,fps=15,quality=8)
     if torch.cuda.is_available(): torch.cuda.empty_cache()
-    return state_dict, video_path
+    return state, vp
 
 @spaces.GPU(duration=120)
-def extract_glb(
-    state_dict: dict,
-    mesh_simplify: float,
-    texture_size: int,
-    req: gr.Request,
-) -> Tuple[str, str]:
+def extract_glb(state_dict: dict, mesh_simplify: float, texture_size: int, req: gr.Request):
     gs, mesh = unpack_state(state_dict)
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=float(mesh_simplify), texture_size=int(texture_size), verbose=True)
-    glb_path = os.path.join(user_dir, 'sample.glb')
-    glb.export(glb_path)
+    ud = os.path.join(TMP_DIR, str(req.session_hash)); os.makedirs(ud, exist_ok=True)
+    glb = postprocessing_utils.to_glb(gs,mesh,simplify=mesh_simplify,texture_size=texture_size,verbose=True)
+    gp = os.path.join(ud,'sample.glb'); glb.export(gp)
     if torch.cuda.is_available(): torch.cuda.empty_cache()
-    return glb_path, glb_path
+    return gp, gp
 
 @spaces.GPU
-def extract_gaussian(
-    state_dict: dict,
-    req: gr.Request
-) -> Tuple[str, str]:
+def extract_gaussian(state_dict: dict, req: gr.Request):
     gs, _ = unpack_state(state_dict)
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-    gaussian_path = os.path.join(user_dir, 'sample.ply')
-    gs.save_ply(gaussian_path)
+    ud = os.path.join(TMP_DIR, str(req.session_hash)); os.makedirs(ud, exist_ok=True)
+    pp = os.path.join(ud,'sample.ply'); gs.save_ply(pp)
     if torch.cuda.is_available(): torch.cuda.empty_cache()
-    return gaussian_path, gaussian_path
+    return pp, pp
 
-# --- Gradio UI Definition ---
-with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
+# --- Gradio UI ---
+with gr.Blocks(delete_cache=(600,600), title="TRELLIS Text-to-3D") as demo:
     gr.Markdown("""
-    # Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
+    # Text to 3D Asset with TRELLIS
     """)
-
-    # State buffer
     output_buf = gr.State()
-
     with gr.Row():
         with gr.Column(scale=1):
             text_prompt = gr.Textbox(label="Text Prompt", lines=5)
-            with gr.Accordion(label="Generation Settings", open=False):
+            with gr.Accordion("Generation Settings", open=False):
                 seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                gr.Markdown("---\n**Stage 1**")
-                ss_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
-                ss_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
-                gr.Markdown("---\n**Stage 2**")
-                slat_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
-                slat_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
-            generate_btn = gr.Button("Generate 3D Preview", variant="primary")
-            with gr.Accordion(label="GLB Extraction Settings", open=True):
-                mesh_simplify = gr.Slider(0.9, 0.99, label="Simplify Factor", value=0.95, step=0.01)
-                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
+                gr.Markdown("--- Stage 1 ---")
+                ss_guidance_strength = gr.Slider(0.0,15.0,label="Guidance Strength",value=7.5,step=0.1)
+                ss_sampling_steps = gr.Slider(10,50,label="Steps",value=25,step=1)
+                gr.Markdown("--- Stage 2 ---")
+                slat_guidance_strength = gr.Slider(0.0,15.0,label="Guidance Strength",value=7.5,step=0.1)
+                slat_sampling_steps = gr.Slider(10,50,label="Steps",value=25,step=1)
+            generate_btn = gr.Button("Generate 3D Preview")
+            with gr.Accordion("GLB Extraction Settings", open=True):
+                mesh_simplify = gr.Slider(0.9,0.99,label="Simplify",value=0.95,step=0.01)
+                texture_size = gr.Slider(512,2048,label="Texture Size",value=1024,step=512)
             extract_glb_btn = gr.Button("Extract GLB", interactive=False)
-            extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False)
-            download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-            download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False)
+            extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+            download_glb = gr.DownloadButton("Download GLB", interactive=False)
+            download_gs = gr.DownloadButton("Download Gaussian", interactive=False)
         with gr.Column(scale=1):
-            video_output = gr.Video(label="3D Preview", autoplay=True, loop=True)
-            model_output = gr.Model3D(label="Extracted Model Preview")
+            video_output = gr.Video(autoplay=True,loop=True)
+            model_output = gr.Model3D()
 
-    # --- Event handlers ---
-    demo.load(start_session)  # FIX: remove inputs/outputs kwargs
-    demo.unload(end_session)  # FIX: remove inputs/outputs kwargs
+    # --- Handlers ---
+    demo.load(start_session)
+    demo.unload(end_session)
 
-    # Align indentation to one level under Blocks
     generate_event = generate_btn.click(
         get_seed,
-        inputs=[randomize_seed, seed],
-        outputs=[seed],
+        inputs=[randomize_seed,seed], outputs=[seed]
     ).then(
         text_to_3d,
-        inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
-        outputs=[output_buf, video_output],
-    ).then(
-        lambda: (extract_glb_btn.update(interactive=True), extract_gs_btn.update(interactive=True)),
-        outputs=[extract_glb_btn, extract_gs_btn],
-    )
+        inputs=[text_prompt,seed,ss_guidance_strength,ss_sampling_steps,slat_guidance_strength,slat_sampling_steps],
+        outputs=[output_buf,video_output]
+    ).then(lambda: (extract_glb_btn.update(interactive=True),extract_gs_btn.update(interactive=True)), outputs=[extract_glb_btn,extract_gs_btn])
 
-    extract_glb_event = extract_glb_btn.click(
+    extract_glb_btn.click(
         extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, download_glb],
-    ).then(
-        lambda: download_glb.update(interactive=True),
-        outputs=[download_glb],
-    )
+        inputs=[output_buf,mesh_simplify,texture_size],
+        outputs=[model_output,download_glb]
+    ).then(lambda: download_glb.update(interactive=True), outputs=[download_glb])
 
-    extract_gs_event = extract_gs_btn.click(
+    extract_gs_btn.click(
         extract_gaussian,
-        inputs=[output_buf],
-        outputs=[model_output, download_gs],
-    ).then(
-        lambda: download_gaussian.update(interactive=True),
-        outputs=[download_gs],
-    )
+        inputs=[output_buf], outputs=[model_output,download_gs]
+    ).then(lambda: download_gs.update(interactive=True), outputs=[download_gs])
 
-    # Clear callbacks
-    model_output.clear(
-        lambda: (download_glb.update(interactive=False), download_gs.update(interactive=False)),
-        outputs=[download_glb, download_gs],
-    )
-    video_output.clear(
-        lambda: (extract_glb_btn.update(interactive=False), extract_gs_btn.update(interactive=False), download_glb.update(interactive=False), download_gs.update(interactive=False)),
-        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
-    )
+    model_output.clear(lambda: (download_glb.update(interactive=False),download_gs.update(interactive=False)), outputs=[download_glb,download_gs])
+    video_output.clear(lambda: (extract_glb_btn.update(interactive=False),extract_gs_btn.update(interactive=False),download_glb.update(interactive=False),download_gs.update(interactive=False)), outputs=[extract_glb_btn,extract_gs_btn,download_glb,download_gs])
 
 if __name__ == "__main__":
-    # Removed torch_dtype argument to match current API
-    pipeline = TrellisTextTo3DPipeline.from_pretrained(
-        "JeffreyXiang/TRELLIS-text-xlarge"
-    )
-    if torch.cuda.is_available(): pipeline = pipeline.to("cuda")
     demo.queue().launch(debug=True)