Update app.py

app.py

@@ -13,39 +13,29 @@ from diffusers import StableVideoDiffusionPipeline, StableVideoDragNUWAPipeline
 from diffusers.utils import export_to_video, export_to_gif
 from PIL import Image
 
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
 fps25Pipe = StableVideoDiffusionPipeline.from_pretrained(
     "vdo/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float16, variant="fp16"
 )
-fps25Pipe.to("cuda")
+fps25Pipe.to(device)
 
 fps14Pipe = StableVideoDiffusionPipeline.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
 )
-fps14Pipe.to("cuda")
+fps14Pipe.to(device)
 
 dragnuwaPipe = StableVideoDragNUWAPipeline.from_pretrained(
     "a-r-r-o-w/dragnuwa-svd", torch_dtype=torch.float16, variant="fp16", low_cpu_mem_usage=False, device_map=None
 )
-dragnuwaPipe.to("cuda")
+dragnuwaPipe.to(device)
 
 max_64_bit_int = 2**63 - 1
 
-def animate(
-    image: Image,
-    seed: Optional[int] = 42,
-    randomize_seed: bool = True,
-    motion_bucket_id: int = 127,
-    fps_id: int = 25,
-    noise_aug_strength: float = 0.1,
-    decoding_t: int = 3,
-    video_format: str = "mp4",
-    frame_format: str = "webp",
-    version: str = "auto",
-    width: int = 1024,
-    height: int = 576,
-    motion_control: bool = False,
-    num_inference_steps: int = 25
-):
+def animate(image, seed=42, randomize_seed=True, motion_bucket_id=127, fps_id=25, noise_aug_strength=0.1, decoding_t=3, video_format="mp4", frame_format="webp", version="auto", width=1024, height=576, motion_control=False, num_inference_steps=25):
     start = time.time()
 
     if image is None:
@@ -108,36 +98,18 @@ def animate(
     str(secondes) + " sec."
     
     return [
-        # Display for video
-        gr.update(value = video_path, visible = video_format != "gif"),
-        # Display for gif
-        gr.update(value = gif_path, visible = video_format == "gif"),
-        # Download button
-        gr.update(label = "💾 Download animation in *." + video_format + " format", value=result_path, visible=True),
-        # Frames
-        gr.update(label = "Generated frames in *." + frame_format + " format", format = frame_format, value = frames, visible = True),
-        # Used seed
+        gr.update(value=video_path, visible=video_format != "gif"),
+        gr.update(value=gif_path, visible=video_format == "gif"),
+        gr.update(label="💾 Download animation in *." + video_format + " format", value=result_path, visible=True),
+        gr.update(label="Generated frames in *." + frame_format + " format", format=frame_format, value=frames, visible=True),
         seed,
-        # Information
-        gr.update(value = information, visible = True),
-        # Reset button
-        gr.update(visible = True)
+        gr.update(value=information, visible=True),
+        gr.update(visible=True)
     ]
 
 @torch.no_grad()
 @spaces.GPU(queue=False)
-def animate_on_gpu(
-    image_data: Union[Image.Image, List[Image.Image]],
-    seed: Optional[int] = 42,
-    motion_bucket_id: int = 127,
-    fps_id: int = 6,
-    noise_aug_strength: float = 0.1,
-    decoding_t: int = 3,
-    version: str = "svdxt",
-    width: int = 1024,
-    height: int = 576,
-    num_inference_steps: int = 25
-):
+def animate_on_gpu(image_data, seed=42, motion_bucket_id=127, fps_id=6, noise_aug_strength=0.1, decoding_t=3, version="svdxt", width=1024, height=576, num_inference_steps=25):
     generator = torch.manual_seed(seed)
 
     if version == "dragnuwa":
@@ -149,37 +121,29 @@ def animate_on_gpu(
 
 
 def resize_image(image, output_size=(1024, 576)):
-    # Do not touch the image if the size is good
     if image.width == output_size[0] and image.height == output_size[1]:
         return image
 
-    # Calculate aspect ratios
-    target_aspect = output_size[0] / output_size[1]  # Aspect ratio of the desired size
-    image_aspect = image.width / image.height  # Aspect ratio of the original image
+    target_aspect = output_size[0] / output_size[1]
+    image_aspect = image.width / image.height
 
-    # Resize if the original image is larger
     if image_aspect > target_aspect:
-        # Resize the image to match the target height, maintaining aspect ratio
         new_height = output_size[1]
         new_width = int(new_height * image_aspect)
         resized_image = image.resize((new_width, new_height), Image.LANCZOS)
-        # Calculate coordinates for cropping
         left = (new_width - output_size[0]) / 2
         top = 0
         right = (new_width + output_size[0]) / 2
         bottom = output_size[1]
     else:
-        # Resize the image to match the target width, maintaining aspect ratio
         new_width = output_size[0]
         new_height = int(new_width / image_aspect)
         resized_image = image.resize((new_width, new_height), Image.LANCZOS)
-        # Calculate coordinates for cropping
         left = 0
         top = (new_height - output_size[1]) / 2
         right = output_size[0]
         bottom = (new_height + output_size[1]) / 2
 
-    # Crop the image
     return resized_image.crop((left, top, right, bottom))
 
 def reset():
@@ -232,7 +196,7 @@ with gr.Blocks() as demo:
               randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
           generate_btn = gr.Button(value="🚀 Animate", variant="primary")
-          reset_btn = gr.Button(value="🧹 Reinit page", variant="stop", elem_id="reset_button", visible = False)
+          reset_btn = gr.Button(value="🧹 Reinit page", variant="stop", elem_id="reset_button", visible=False)
 
       with gr.Column():
           video_output = gr.Video(label="Generated video", format="mp4", autoplay=True, show_download_button=False)
@@ -266,7 +230,7 @@ with gr.Blocks() as demo:
       reset_btn
   ], api_name="video")
 
-  reset_btn.click(fn = reset, inputs = [], outputs = [
+  reset_btn.click(fn=reset, inputs=[], outputs=[
       image,
       seed,
       randomize_seed,
@@ -281,7 +245,7 @@ with gr.Blocks() as demo:
       height,
       motion_control,
       num_inference_steps
-  ], queue = False, show_progress = False)
+  ], queue=False, show_progress=False)
     
   gr.Examples(
     examples=[