spacesgpu

app.py

@@ -18,7 +18,6 @@ from PIL import Image
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
-from contextlib import contextmanager
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 # Constants
@@ -84,23 +83,6 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
 def get_seed(randomize_seed: bool, seed: int) -> int:
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
-
-# Example class-based or function-based context manager
-@contextmanager
-def pipeline_on_gpu(pipeline, device="cuda"):
-    """
-    Context manager that places the pipeline on GPU at enter,
-    then on exit puts it to CPU to free VRAM.
-    """
-    # Move pipeline from CPU to GPU (if needed)
-    pipeline.to(device)
-    try:
-        yield pipeline
-    finally:
-        # Move pipeline back to CPU and clear CUDA cache
-        pipeline.to("cpu")
-        torch.cuda.empty_cache()
-
 @spaces.GPU
 def generate_flux_image(
     prompt: str,
@@ -113,25 +95,20 @@ def generate_flux_image(
     lora_scale: float,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ) -> Image.Image:
-    """Generate image using Flux pipeline only on GPU during generation."""
+    """Generate image using Flux pipeline"""
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Use the context manager to keep the pipeline on GPU just while generating.
-    with pipeline_on_gpu(flux_pipeline, device=device) as gpu_pipeline:
-        result = gpu_pipeline(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-        )
-
-    # Once we leave the context manager, the pipeline is moved back to CPU.
-    image = result.images[0]
+    image = flux_pipeline(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
+    ).images[0]
     
     # Save the generated image
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -152,48 +129,28 @@ def image_to_3d(
     slat_sampling_steps: int,
     req: gr.Request,
 ) -> Tuple[dict, str]:
-    # Clear CUDA cache before starting
-    torch.cuda.empty_cache()
-    
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    
-    try:
-        with pipeline_on_gpu(trellis_pipeline, device=device) as gpu_pipeline:
-            outputs = gpu_pipeline.run(
-                image,
-                seed=seed,
-                formats=["gaussian", "mesh"],
-                preprocess_image=False,
-                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,
-                },
-            )
-            
-            # Create video while model is still on GPU
-            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 = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-            
-            # Pack state while tensors are still on GPU
-            state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-            
-    except Exception as e:
-        # Ensure cleanup on error
-        torch.cuda.empty_cache()
-        raise e
-    
-    # Save video after GPU operations are complete
+    outputs = trellis_pipeline.run(
+        image,
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        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,
+        },
+    )
+    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 = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
     video_path = os.path.join(user_dir, 'sample.mp4')
     imageio.mimsave(video_path, video, fps=15)
-    
-    # Final cleanup
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
     torch.cuda.empty_cache()
-    
     return state, video_path
 
 @spaces.GPU(duration=90)