Update gradio_app.py

gradio_app.py

@@ -10,11 +10,7 @@ from attn_ctrl.attention_control import (AttentionStore,
                                          register_temporal_self_attention_control, 
                                          register_temporal_self_attention_flip_control,
 )
-from torch.amp import autocast
-import gc
-
-# Set PYTORCH_CUDA_ALLOC_CONF
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+from torch.cuda.amp import autocast
 
 # Set up device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -32,7 +28,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 +37,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()
@@ -68,7 +64,6 @@ register_temporal_self_attention_flip_control(pipe.unet, controller, controller_
 def cuda_memory_cleanup():
     torch.cuda.empty_cache()
     torch.cuda.ipc_collect()
-    gc.collect()
 
 def check_outputs_folder(folder_path):
     if os.path.exists(folder_path) and os.path.isdir(folder_path):
@@ -87,51 +82,47 @@ def check_outputs_folder(folder_path):
 @torch.no_grad()
 def infer(frame1_path, frame2_path):
     seed = 42
-    num_inference_steps = 5  # Reduced from 10
+    num_inference_steps = 10
     noise_injection_steps = 0
     noise_injection_ratio = 0.5
     weighted_average = False
+    decode_chunk_size = 8
 
     generator = torch.Generator(device)
     if seed is not None:
         generator = generator.manual_seed(seed)
 
     frame1 = load_image(frame1_path)
-    frame1 = frame1.resize((256, 144))  # Reduced from (512, 288)
+    frame1 = frame1.resize((512, 288))
 
     frame2 = load_image(frame2_path)
-    frame2 = frame2.resize((256, 144))  # Reduced from (512, 288)
+    frame2 = frame2.resize((512, 288))
 
-    # Clear CUDA cache
     cuda_memory_cleanup()
 
-    try:
-        with autocast():
-            frames = pipe(
-                image1=frame1, 
-                image2=frame2, 
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                weighted_average=weighted_average,
-                noise_injection_steps=noise_injection_steps,
-                noise_injection_ratio=noise_injection_ratio,
-            ).frames[0]
-
-        frames = [frame.cpu() for frame in frames]
-
-        out_dir = "result"
-        check_outputs_folder(out_dir)
-        os.makedirs(out_dir, exist_ok=True)
-        out_path = "result/video_result.gif"
-
-        return "done"
-    except RuntimeError as e:
-        if "CUDA out of memory" in str(e):
-            return "Error: CUDA out of memory. Try reducing the image size or using fewer inference steps."
-        else:
-            return f"An error occurred: {str(e)}"
-    finally:
-        cuda_memory_cleanup()
+    with autocast():
+        frames = pipe(image1=frame1, image2=frame2, 
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            weighted_average=weighted_average,
+            noise_injection_steps=noise_injection_steps,
+            noise_injection_ratio=noise_injection_ratio,
+            decode_chunk_size=decode_chunk_size
+        ).frames[0]
+
+    frames = [frame.cpu() for frame in frames]
+
+    out_dir = "result"
+    check_outputs_folder(out_dir)
+    os.makedirs(out_dir, exist_ok=True)
+    out_path = "result/video_result.gif"
+
+    return "done"
+
+@torch.no_grad()
+def load_model():
+    global pipe
+    pipe = pipe.to(device)
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -151,4 +142,6 @@ with gr.Blocks() as demo:
         show_api=False
     )
 
-demo.queue(max_size=1).launch(show_api=False, show_error=True, share=True)
+    demo.load(load_model)
+
+demo.queue(max_size=1).launch(show_api=False, show_error=True)