Update app.py

app.py

@@ -320,6 +320,16 @@ def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 30, gu
         logger.error(f"Error in gradio_generate: {str(e)}")
         raise
 
+# Conditional GPU decorator
+def conditional_gpu_decorator(duration=None):
+    def decorator(func):
+        if IS_ZEROGPU:
+            return spaces.GPU(duration=duration)(func) if duration else spaces.GPU()(func)
+        return func
+    return decorator
+
+# Always apply @spaces.GPU for ZeroGPU in start_session
+@spaces.GPU() if IS_ZEROGPU else lambda x: x
 def start_session(req: gr.Request):
     try:
         save_dir = os.path.join(TMP_DIR, str(req.session_hash))