Update app.py

app.py

@@ -155,7 +155,7 @@ from einops import rearrange
 from shap_e.diffusion.sample import sample_latents
 from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
 from shap_e.models.download import load_model, load_config
-from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
+from shap_e.util.notebooks import create_pan_cameras
 from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
 import math
 import time
@@ -171,6 +171,29 @@ from src.utils.camera_util import (
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground
 
+def decode_latent_images(
+    xm: Union[Transmitter, VectorDecoder],
+    latent: torch.Tensor,
+    cameras: DifferentiableCameraBatch,
+    rendering_mode: str = "stf",
+    params = None,
+    background_color: torch.Tensor = torch.tensor([255.0, 255.0, 255.0], dtype=torch.float32),
+):
+    params = params if params is not None else (xm.encoder if isinstance(xm, Transmitter) else xm).bottleneck_to_params(
+            latent[None]
+        )
+    params = xm.renderer.update(params)
+    decoded = xm.renderer.render_views(
+        AttrDict(cameras=cameras),
+        params=params,
+        options=AttrDict(rendering_mode=rendering_mode, render_with_direction=False),
+    )
+    bg_color = background_color.to(decoded.channels.device)
+    images = bg_color * decoded.transmittance + (1 - decoded.transmittance) * decoded.channels
+
+    # arr = decoded.channels.clamp(0, 255).to(torch.uint8)[0].cpu().numpy()
+    return images
+
 def create_custom_cameras(size: int, device: torch.device, azimuths: list, elevations: list, 
                           fov_degrees: float, distance: float) -> DifferentiableCameraBatch:
     # Object is in a 2x2x2 bounding box (-1 to 1 in each dimension)
@@ -654,7 +677,7 @@ def create_demo():
             )
 
         # Set up event handlers
-        @spaces.GPU(duration=90)  # Reduced duration to 20 seconds
+        @spaces.GPU(duration=20)  # Reduced duration to 20 seconds
         def generate(prompt, guidance_scale, num_steps):
             try:
                 # Ensure PyTorch3D works with CUDA