Update app.py

app.py

@@ -15,10 +15,14 @@ import random
 token = os.environ['HF_TOKEN']
 login(token=token)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+torch.cuda.max_memory_allocated(device=device)
+torch.cuda.empty_cache()
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt-1-1")
 #pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-pipe = pipe.to(device)
+
 pipe.enable_xformers_memory_efficient_attention()
+pipe = pipe.to(device)
+torch.cuda.empty_cache()
 
 max_64_bit_int = 2**63 - 1
 
@@ -40,7 +44,7 @@ def sample(
     if(randomize_seed):
         seed = random.randint(0, max_64_bit_int)
     generator = torch.manual_seed(seed)
-    
+    torch.cuda.empty_cache()
     os.makedirs(output_folder, exist_ok=True)
     base_count = len(glob(os.path.join(output_folder, "*.mp4")))
     video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
@@ -48,7 +52,7 @@ def sample(
     frames = pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=0.1, num_frames=25).frames[0]
     export_to_video(frames, video_path, fps=fps_id)
     torch.manual_seed(seed)
-    
+    torch.cuda.empty_cache()
     return video_path, seed
 
 def resize_image(image, output_size=(1024, 578)):
@@ -80,6 +84,7 @@ def resize_image(image, output_size=(1024, 578)):
 
     # Crop the image
     cropped_image = resized_image.crop((left, top, right, bottom))
+    torch.cuda.empty_cache()
     return cropped_image
 
 with gr.Blocks() as demo: