Update gradio_app.py

gradio_app.py

@@ -10,7 +10,7 @@ from attn_ctrl.attention_control import (AttentionStore,
                                          register_temporal_self_attention_control, 
                                          register_temporal_self_attention_flip_control,
 )
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 import gc
 
 # Set PYTORCH_CUDA_ALLOC_CONF
@@ -32,7 +32,7 @@ pipe = FrameInterpolationWithNoiseInjectionPipeline.from_pretrained(
     scheduler=noise_scheduler,
     variant="fp16",
     torch_dtype=torch.float16, 
-)
+).to(device)
 ref_unet = pipe.ori_unet
 
 # Compute delta w
@@ -41,14 +41,14 @@ finetuned_unet = UNetSpatioTemporalConditionModel.from_pretrained(
     checkpoint_dir,
     subfolder="unet",
     torch_dtype=torch.float16,
-) 
+).to(device)
 assert finetuned_unet.config.num_frames == 14
 ori_unet = UNetSpatioTemporalConditionModel.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid",
     subfolder="unet",
     variant='fp16',
     torch_dtype=torch.float16,
-)
+).to(device)
 
 finetuned_state_dict = finetuned_unet.state_dict()
 ori_state_dict = ori_unet.state_dict()
@@ -105,15 +105,8 @@ def infer(frame1_path, frame2_path):
     # Clear CUDA cache
     cuda_memory_cleanup()
 
-    # Move model to CPU and clear CUDA cache
-    pipe.to("cpu")
-    cuda_memory_cleanup()
-
-    # Move model back to GPU
-    pipe.to(device)
-
     try:
-        with autocast(device_type='cuda', dtype=torch.float16):
+        with autocast():
             frames = pipe(
                 image1=frame1, 
                 image2=frame2, 
@@ -138,15 +131,8 @@ def infer(frame1_path, frame2_path):
         else:
             return f"An error occurred: {str(e)}"
     finally:
-        # Move model back to CPU and clear CUDA cache
-        pipe.to("cpu")
         cuda_memory_cleanup()
 
-@torch.no_grad()
-def load_model():
-    global pipe
-    pipe = pipe.to(device)
-
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# Keyframe Interpolation with Stable Video Diffusion")
@@ -165,6 +151,4 @@ with gr.Blocks() as demo:
         show_api=False
     )
 
-    demo.load(load_model)
-
-demo.queue(max_size=1).launch(show_api=False, show_error=True)
+demo.queue(max_size=1).launch(show_api=False, show_error=True, share=True)