Update app_gradio.py

app_gradio.py

@@ -24,6 +24,7 @@ from invert_utils import ddim_inversion as dd_inversion
 from gifs_filter import filter
 import subprocess
 import uuid
+import tempfile
 
 from huggingface_hub import snapshot_download
 
@@ -159,34 +160,44 @@ def generate_output(image, apply_filter, prompt: str, num_seeds: int = 3, lambda
     if prompt is None:
         raise gr.Error("You forgot to describe the motion !")
     """Main function to generate output GIFs"""
-    unique_id = str(uuid.uuid4())
-    exp_dir = f"static/app_tmp_{unique_id}"
+    
+    # Create a temporary directory for this session
+    exp_dir = tempfile.mkdtemp(prefix="app_tmp_")
     os.makedirs(exp_dir, exist_ok=True)
     
     # Save the input image temporarily
-    temp_image_path = os.path.join(exp_dir, "temp_input.png")
+    unique_id = str(uuid.uuid4())
+    temp_image_path = os.path.join(exp_dir, f"temp_input_{unique_id}.png")
 
     image = Image.open(image)
     image = image.resize((256, 256), Image.LANCZOS)
 
     image.save(temp_image_path)
     
-    # Generate the GIFs
-    generated_gifs = process_video(
-        num_frames=10,
-        num_seeds=num_seeds,
-        generator=None,
-        exp_dir=exp_dir,
-        load_name=temp_image_path,
-        caption=prompt,
-        lambda_=1 - lambda_value
-    )
-
-    if apply_filter == True:
-        print("APPLYING FILTER")
-        # Apply filtering (assuming filter function is imported)
-        filtered_gifs = filter(generated_gifs, temp_image_path)   
-        return filtered_gifs, filtered_gifs
+    try:
+        # Attempt to process video
+        generated_gifs = process_video(
+            num_frames=10,
+            num_seeds=num_seeds,
+            generator=None,
+            exp_dir=exp_dir,
+            load_name=temp_image_path,
+            caption=prompt,
+            lambda_=1 - lambda_value
+        )
+    except Exception as e:
+        torch.cuda.empty_cache()  # Clear CUDA cache in case of failure
+        raise gr.Error(f"Video processing failed: {str(e)}") from e
+    
+    if apply_filter:
+        try:
+            print("APPLYING FILTER")
+            # Attempt to apply filtering
+            filtered_gifs = filter(generated_gifs, temp_image_path)
+            return filtered_gifs, filtered_gifs
+        except Exception as e:
+            torch.cuda.empty_cache()  # Clear CUDA cache in case of failure
+            raise gr.Error(f"Filtering failed: {str(e)}") from e
     else:
         print("NOT APPLYING FILTER")
         return generated_gifs, generated_gifs