Reduce the usage of GPU (#20)

- Reduce the usage of GPU (d76675441b3ffd721192a61fdeb81cbb31fc9f6d)


Co-authored-by: Fabrice TIERCELIN <Fabrice-TIERCELIN@users.noreply.huggingface.co>

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import torch
 import os
+import random
+import spaces
 from glob import glob
 from pathlib import Path
 from typing import Optional
@@ -9,9 +11,6 @@ from diffusers import StableVideoDiffusionPipeline
 from diffusers.utils import export_to_video
 from PIL import Image
 
-import random
-import spaces
-
 fps25Pipe = StableVideoDiffusionPipeline.from_pretrained(
     "vdo/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float16, variant="fp16"
 )
@@ -24,8 +23,7 @@ fps14Pipe.to("cuda")
 
 max_64_bit_int = 2**63 - 1
 
-@spaces.GPU(duration=120)
-def sample(
+def animate(
     image: Image,
     seed: Optional[int] = 42,
     randomize_seed: bool = True,
@@ -35,7 +33,6 @@ def sample(
     decoding_t: int = 3,
     frame_format: str = "webp",
     version: str = "auto",
-    device: str = "cuda",
     output_folder: str = "outputs",
 ):
     if image.mode == "RGBA":
@@ -43,20 +40,47 @@ def sample(
         
     if randomize_seed:
         seed = random.randint(0, max_64_bit_int)
-    generator = torch.manual_seed(seed)
+
+    frames = animate_on_gpu(
+        image,
+        seed,
+        randomize_seed,
+        motion_bucket_id,
+        fps_id,
+        noise_aug_strength,
+        decoding_t,
+        frame_format,
+        version
+    )
     
     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")
 
-    if version == "svdxt" or (14 < fps_id and version != "svd"):
-        frames = fps25Pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25).frames[0]
-    else:
-        frames = fps14Pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25).frames[0]
     export_to_video(frames, video_path, fps=fps_id)
     
     return video_path, gr.update(value=video_path, visible=True), gr.update(label="Generated frames in *." + frame_format + " format", format = frame_format, value = frames, visible=True), seed
 
+@spaces.GPU(duration=120)
+def animate_on_gpu(
+    image: Image,
+    seed: Optional[int] = 42,
+    randomize_seed: bool = True,
+    motion_bucket_id: int = 127,
+    fps_id: int = 6,
+    noise_aug_strength: float = 0.1,
+    decoding_t: int = 3,
+    frame_format: str = "webp",
+    version: str = "auto"
+):
+    generator = torch.manual_seed(seed)
+    
+    if version == "svdxt" or (14 < fps_id and version != "svd"):
+        return fps25Pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25).frames[0]
+    else:
+        return fps14Pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25).frames[0]
+
+
 def resize_image(image, output_size=(1024, 576)):
     # Calculate aspect ratios
     target_aspect = output_size[0] / output_size[1]  # Aspect ratio of the desired size
@@ -117,7 +141,7 @@ with gr.Blocks() as demo:
           gallery = gr.Gallery(label="Generated frames", visible=False)
       
   image.upload(fn=resize_image, inputs=image, outputs=image, queue=False)
-  generate_btn.click(fn=sample, inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id, noise_aug_strength, decoding_t, frame_format, version], outputs=[video, download_button, gallery, seed], api_name="video")
+  generate_btn.click(fn=animate, inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id, noise_aug_strength, decoding_t, frame_format, version], outputs=[video, download_button, gallery, seed], api_name="video")
     
   gr.Examples(
     examples=[
@@ -127,7 +151,7 @@ with gr.Blocks() as demo:
     ],
     inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id, noise_aug_strength, decoding_t, frame_format, version],
     outputs=[video, download_button, gallery, seed],
-    fn=sample,
+    fn=animate,
     run_on_click=True,
     cache_examples=False,
   )