# @title Do the Run! markdown Generation takes about 10 mins for `svd_xt` on T4 (Colab free plan). Please be patient...
import gradio as gr
import random


def infer(input_path: str, resize_image: bool, n_frames: int, n_steps: int, seed: str, decoding_t: int, fps_id: int, motion_bucket_id: int, cond_aug: float, skip_filter: bool = False) -> str:
  if seed == "random":
    seed = random.randint(0, 2**32)
  seed = int(seed)
  output_paths = sample(
    input_path=input_path,
    resize_image=resize_image,
    num_frames=n_frames,
    num_steps=n_steps,
    fps_id=fps_id,
    motion_bucket_id=motion_bucket_id,
    cond_aug=cond_aug,
    seed=seed,
    decoding_t=decoding_t,  # Number of frames decoded at a time! This eats most VRAM. Reduce if necessary.
    device=device,
    skip_filter=skip_filter,
  )
  return output_paths[0]


with gr.Blocks() as demo:
  with gr.Column():
    image = gr.Image(label="input image", type="filepath")
    resize_image = gr.Checkbox(label="resize to optimal size", value=True)
    btn = gr.Button("Run")
    with gr.Accordion(label="Advanced options", open=False):
      n_frames = gr.Number(precision=0, label="number of frames", value=num_frames)
      n_steps = gr.Number(precision=0, label="number of steps", value=num_steps)
      seed = gr.Text(value="random", label="seed (integer or 'random')",)
      decoding_t = gr.Number(precision=0, label="number of frames decoded at a time", value=2)
      fps_id = gr.Number(precision=0, label="frames per second", value=6)
      motion_bucket_id = gr.Number(precision=0, value=127, label="motion bucket id")
      cond_aug = gr.Number(label="condition augmentation factor", value=0.02)
      skip_filter = gr.Checkbox(value=False, label="skip nsfw/watermark filter")
  with gr.Column():
    video_out = gr.Video(label="generated video")
  examples = [
      ["https://user-images.githubusercontent.com/33302880/284758167-367a25d8-8d7b-42d3-8391-6d82813c7b0f.png"]
  ]
  inputs = [image, resize_image, n_frames, n_steps, seed, decoding_t, fps_id, motion_bucket_id, cond_aug, skip_filter]
  outputs = [video_out]
  btn.click(infer, inputs=inputs, outputs=outputs)
  gr.Examples(examples=examples, inputs=inputs, outputs=outputs, fn=infer)
  demo.queue().launch(debug=True, share=True, show_error=True)