Update app.py (#8)

- Update app.py (9760de5dfe36958a8635b547ea94c77632ccca96)
- Update app.py (9552bee688ac9e74e474b637cc1584a4f4b2c1f9)


Co-authored-by: Linoy Tsaban <[email protected]>

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,25 @@ 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
+        
+        # Clean up the session directory
+        session_dir = os.path.join("/tmp/gradio", session_id)
+        if os.path.exists(session_dir):
+            shutil.rmtree(session_dir)
+            print(f"Cleaned up session directory: {session_dir}")
+            
+    except Exception as e:
+        print(f"Error during session cleanup: {e}")
+
 @spaces.GPU()
 def read_video(video) -> torch.Tensor:
     """
@@ -136,7 +155,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 +305,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 +316,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 +469,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()