Update app.py

app.py

@@ -322,23 +322,22 @@ class ShapERenderer:
             guidance_scale = float(guidance_scale)
             
             with torch.amp.autocast('cuda'):  # Use automatic mixed precision
-                # spaces duration is 20 seconds, so we need to be careful here
-            
+                # Generate latents directly without nested spaces.GPU context
                 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,
-                use_karras=True,
-                karras_steps=num_steps,
-                sigma_min=1e-3,
-                sigma_max=160,
-                s_churn=0,
-            )
+                    progress=True,
+                    clip_denoised=True,
+                    use_fp16=True,
+                    use_karras=True,
+                    karras_steps=num_steps,
+                    sigma_min=1e-3,
+                    sigma_max=160,
+                    s_churn=0,
+                )
 
             # Render the 6 views we need with specific viewing angles
             size = 320  # Size of each rendered image
@@ -539,6 +538,7 @@ def create_demo():
             )
 
         # Set up event handlers
+        @spaces.GPU(duration=20)  # Add GPU decorator to the generate function
         def generate(prompt, guidance_scale, num_steps):
             try:
                 torch.cuda.empty_cache()  # Clear GPU memory before starting
@@ -551,7 +551,7 @@ def create_demo():
                 print(error_msg)
                 return None, error_msg
 
-        @spaces.GPU(duration=20)  # Reduced duration to 20 seconds
+        @spaces.GPU(duration=20)
         def refine(input_image, prompt, steps, guidance_scale):
             try:
                 torch.cuda.empty_cache()  # Clear GPU memory before starting