Update app.py

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 
 
 def tune_video_predict(
+    pipe_id: str,
     prompt: str,
     video_length: int,
     height: int,
@@ -13,7 +14,7 @@ def tune_video_predict(
     num_inference_steps: int,
     guidance_scale: float,
 ):
-    unet = UNet3DConditionModel.from_pretrained('Tune-A-Video-library/a-man-is-surfing', subfolder='unet', torch_dtype=torch.float16).to('cuda')
+    unet = UNet3DConditionModel.from_pretrained(pipe_id, subfolder='unet', torch_dtype=torch.float16).to('cuda')
     pipe = TuneAVideoPipeline.from_pretrained('CompVis/stable-diffusion-v1-4', unet=unet, torch_dtype=torch.float16).to("cuda")
     video = pipe(prompt, video_length=video_length, height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).videos
     output_path = save_videos_grid(video, save_path='output', path=f"{prompt}.gif")
@@ -21,6 +22,13 @@ def tune_video_predict(
 
 
 demo_inputs = [
+    gr.inputs.Dropdown(
+        label="Model",
+        choices=[
+            "Tune-A-Video-library/a-man-is-surfing",
+            "sd-dreambooth-library/mr-potato-head",
+        ]
+    ),
     gr.inputs.Textbox(
         label="Prompt",
         default='a flower blooming'
@@ -67,8 +75,10 @@ demo_inputs = [
 demo_outputs = gr.outputs.Video(type="gif", label="Output")
 
 examples = [
-    ["a panda is surfing", 8, 416, 416, 50, 7.5],
-    ["a flower blooming", 5, 416, 416, 50, 7.5],
+    ["Tune-A-Video-library/a-man-is-surfing", "a panda is surfing", 5, 416, 416, 50, 7.5],
+    ["Tune-A-Video-library/a-man-is-surfing", "a flower blooming", 5, 416, 416, 50, 7.5],
+    ["sd-dreambooth-library/mr-potato-head", "sks mr potato head, wearing a pink hat, is surfing.", 5, 416, 416, 50, 7.5],
+    ["sd-dreambooth-library/mr-potato-head", "sks mr potato head is surfing in the forest.", 5, 416, 416, 50, 7.5],
 ]
     
 description = "This generates video from an input text, using [one-shot tuning of diffusion models](https://arxiv.org/abs/2212.11565). To use it, simply input a text."