change to GPU version

README.md

@@ -4,7 +4,7 @@ emoji: 💊
 colorFrom: purple
 colorTo: blue
 sdk: gradio
-sdk_version: 3.38.0
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import numpy as np
 import torch
 import cv2
@@ -113,6 +114,7 @@ def images_to_video(image_list, output_path, fps=10):
     writer.close()
 
 
+@spaces.GPU
 def NaRCan_make_video(edit_canonical, pth_path, frames_path):
     # load NaRCan model
     checkpoint_g_old = torch.load(os.path.join(pth_path, "homography_g.pth"))
@@ -190,6 +192,7 @@ def NaRCan_make_video(edit_canonical, pth_path, frames_path):
     return edit_video_path
 
 
+@spaces.GPU
 def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt, control_type="Lineart"):
     video_name = input_video.split('/')[-1]
     if video_name in video_to_image:
@@ -201,9 +204,9 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
 
     if control_type == "Lineart":
         # Load the control net model for lineart
-        controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_lineart")
+        controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_lineart", torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", controlnet=controlnet
+            "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
         )
         pipe.to(device)
         # lineart
@@ -228,9 +231,9 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
     
     else:
         # Load the control net model for canny
-        controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny")
+        controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", controlnet=controlnet
+            "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
         )
         pipe.to(device)
         # canny
@@ -285,10 +288,10 @@ with gr.Blocks(css="style.css") as demo:
     do_inversion = gr.State(value=True)
 
     with gr.Row():
-        input_video = gr.Video(label="Input Video", interactive=False, elem_id="input_video", value='examples/bear.mp4')
-        output_video = gr.Video(label="Edited Video", interactive=False, elem_id="output_video")
-        input_video.style(height=365, width=365)
-        output_video.style(height=365, width=365)
+        input_video = gr.Video(label="Input Video", interactive=False, elem_id="input_video", value='examples/bear.mp4', height=365, width=365)
+        output_video = gr.Video(label="Edited Video", interactive=False, elem_id="output_video", height=365, width=365)
+        # input_video.style(height=365, width=365)
+        # output_video.style(height=365, width=365)
 
 
     with gr.Row():

requirements.txt

@@ -1,5 +1,6 @@
-gradio==3.38.0
+gradio
 numpy
+spaces
 torch==2.3.1
 opencv-python==4.10.0.82
 imageio==2.34.1