Update app.py

app.py

@@ -1,9 +1,11 @@
 import os
+from typing import Union
 # this is a HF Spaces specific hack for ZeroGPU
 import spaces
 
 import sys
 import torch
+from shap_e.models.transmitter.base import Transmitter, VectorDecoder
 
 import torch
 import torch.nn as nn
@@ -319,12 +321,14 @@ class ShapERenderer:
             guidance_scale = float(guidance_scale)
             
             with torch.amp.autocast('cuda'):  # Use automatic mixed precision
-                latents = sample_latents(
-                    batch_size=batch_size,
-                    model=self.model,
-                    diffusion=self.diffusion,
-                    guidance_scale=guidance_scale,
-                    model_kwargs=dict(texts=[prompt] * batch_size),
+                # spaces duration is 20 seconds, so we need to be careful here
+                with spaces.GPU(duration=20):   
+                    latents = sample_latents(
+                        batch_size=batch_size,
+                        model=self.model,
+                        diffusion=self.diffusion,
+                        guidance_scale=guidance_scale,
+                        model_kwargs=dict(texts=[prompt] * batch_size),
                     progress=True,
                     clip_denoised=True,
                     use_fp16=True,
@@ -534,7 +538,6 @@ def create_demo():
             )
 
         # Set up event handlers
-        @spaces.GPU(duration=20)  # Reduced duration to 20 seconds
         def generate(prompt, guidance_scale, num_steps):
             try:
                 torch.cuda.empty_cache()  # Clear GPU memory before starting