Update app.py

app.py

@@ -13,6 +13,9 @@ import imageio
 from controlnet_aux import CannyDetector
 from PIL import Image
 import cv2
+import shutil
+import glob
+from pathlib import Path
 
 FPS = 24
 dtype = torch.bfloat16 
@@ -52,9 +55,42 @@ pipeline.load_lora_weights(
         )
 pipeline.set_adapters([CONTROL_LORAS["canny"]["adapter_name"]], adapter_weights=[1.0])
 
-
 canny_processor = CannyDetector()
 
+current_dir = Path(__file__).parent
+
+def cleanup_session_files(request: gr.Request):
+    """Clean up session-specific temporary files when user disconnects"""
+    try:
+        # Get the session ID from the request
+        session_id = request.session_hash
+        
+        # Check common gradio temp locations
+        gradio_temp_dirs = [
+            "/tmp/gradio",
+            os.path.join(tempfile.gettempdir(), "gradio"),
+            os.path.join(os.getcwd(), "gradio")
+        ]
+        
+        files_cleaned = 0
+        
+        # Clean session directories in gradio temp folders
+        for temp_dir in gradio_temp_dirs:
+            session_dir = os.path.join(temp_dir, session_id)
+            if os.path.exists(session_dir):
+                try:
+                    shutil.rmtree(session_dir)
+                    files_cleaned += 1
+                    print(f"Cleaned up session directory: {session_dir}")
+                except (OSError, FileNotFoundError):
+                    pass
+        
+        if files_cleaned > 0:
+            print(f"Session cleanup completed: {files_cleaned} directories removed for session {session_id}")
+            
+    except Exception as e:
+        print(f"Error during session cleanup: {e}")
+
 @spaces.GPU()
 def read_video(video) -> torch.Tensor:
     """
@@ -136,7 +172,6 @@ def process_input_video(reference_video, width, height,
         # Process video for canny edges
         processed_video = process_video_for_canny(video, width, height)
         
-        # Create a preview video file for display
         output_filename = "control.mp4"
         video_uint8_frames = [(frame * 255).astype(np.uint8) for frame in processed_video]
         with imageio.get_writer(output_filename, fps=FPS, quality=8, macro_block_size=1) as writer: 
@@ -287,7 +322,7 @@ def generate_video(
         return None, seed
 
 # Create Gradio interface
-with gr.Blocks(theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend Deca"), "sans-serif"])) as demo:
+with gr.Blocks(theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend Deca"), "sans-serif"]), delete_cache=(60, 3600)) as demo:
     gr.Markdown(
         """
         # Canny Control LTX Video Distilled
@@ -298,9 +333,6 @@ with gr.Blocks(theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend Deca"),
         """
     )
     
-    # State variables
-    #current_lora_state = gr.State(value=None)
-    
     with gr.Row():
         with gr.Column(scale=1):
     
@@ -454,5 +486,8 @@ with gr.Blocks(theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend Deca"),
         show_progress=True
     )
 
+    # Add the unload event handler for cleanup
+    demo.unload(cleanup_session_files)
+
 if __name__ == "__main__":
     demo.launch()