|
|
|
|
|
from __future__ import annotations |
|
|
|
import os |
|
|
|
import gradio as gr |
|
import torch |
|
|
|
from inference import InferencePipeline |
|
|
|
|
|
class InferenceUtil: |
|
def __init__(self, hf_token: str | None): |
|
self.hf_token = hf_token |
|
|
|
def load_model_info(self, model_id: str) -> tuple[str, str]: |
|
try: |
|
card = InferencePipeline.get_model_card(model_id, self.hf_token) |
|
except Exception: |
|
return "", "" |
|
base_model = getattr(card.data, "base_model", "") |
|
training_prompt = getattr(card.data, "training_prompt", "") |
|
return base_model, training_prompt |
|
|
|
|
|
DESCRIPTION = "# [Tune-A-Video](https://tuneavideo.github.io/)" |
|
if not torch.cuda.is_available(): |
|
DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>" |
|
|
|
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1" |
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN") |
|
pipe = InferencePipeline(HF_TOKEN) |
|
app = InferenceUtil(HF_TOKEN) |
|
|
|
with gr.Blocks(css="style.css") as demo: |
|
gr.Markdown(DESCRIPTION) |
|
gr.DuplicateButton( |
|
value="Duplicate Space for private use", |
|
elem_id="duplicate-button", |
|
visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1", |
|
) |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Group(): |
|
model_id = gr.Dropdown( |
|
label="Model ID", |
|
choices=[ |
|
"Tune-A-Video-library/a-man-is-surfing", |
|
"Tune-A-Video-library/mo-di-bear-guitar", |
|
"Tune-A-Video-library/redshift-man-skiing", |
|
], |
|
value="Tune-A-Video-library/a-man-is-surfing", |
|
) |
|
with gr.Accordion(label="Model info (Base model and prompt used for training)", open=False): |
|
with gr.Row(): |
|
base_model_used_for_training = gr.Text(label="Base model", interactive=False) |
|
prompt_used_for_training = gr.Text(label="Training prompt", interactive=False) |
|
prompt = gr.Textbox(label="Prompt", max_lines=1, placeholder='Example: "A panda is surfing"') |
|
video_length = gr.Slider(label="Video length", minimum=4, maximum=12, step=1, value=8) |
|
fps = gr.Slider(label="FPS", minimum=1, maximum=12, step=1, value=1) |
|
seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, value=0) |
|
with gr.Accordion("Other Parameters", open=False): |
|
num_steps = gr.Slider(label="Number of Steps", minimum=0, maximum=100, step=1, value=50) |
|
guidance_scale = gr.Slider(label="CFG Scale", minimum=0, maximum=50, step=0.1, value=7.5) |
|
|
|
run_button = gr.Button("Generate") |
|
|
|
gr.Markdown( |
|
""" |
|
- It takes a few minutes to download model first. |
|
- Expected time to generate an 8-frame video: 70 seconds with T4, 24 seconds with A10G, (10 seconds with A100) |
|
""" |
|
) |
|
with gr.Column(): |
|
result = gr.Video(label="Result") |
|
with gr.Row(): |
|
examples = [ |
|
[ |
|
"Tune-A-Video-library/a-man-is-surfing", |
|
"A panda is surfing.", |
|
8, |
|
1, |
|
3, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/a-man-is-surfing", |
|
"A racoon is surfing, cartoon style.", |
|
8, |
|
1, |
|
3, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/mo-di-bear-guitar", |
|
"a handsome prince is playing guitar, modern disney style.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/mo-di-bear-guitar", |
|
"a magical princess is playing guitar, modern disney style.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/mo-di-bear-guitar", |
|
"a rabbit is playing guitar, modern disney style.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/mo-di-bear-guitar", |
|
"a baby is playing guitar, modern disney style.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/redshift-man-skiing", |
|
"(redshift style) spider man is skiing.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/redshift-man-skiing", |
|
"(redshift style) black widow is skiing.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/redshift-man-skiing", |
|
"(redshift style) batman is skiing.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
[ |
|
"Tune-A-Video-library/redshift-man-skiing", |
|
"(redshift style) hulk is skiing.", |
|
8, |
|
1, |
|
123, |
|
50, |
|
7.5, |
|
], |
|
] |
|
gr.Examples( |
|
examples=examples, |
|
inputs=[ |
|
model_id, |
|
prompt, |
|
video_length, |
|
fps, |
|
seed, |
|
num_steps, |
|
guidance_scale, |
|
], |
|
outputs=result, |
|
fn=pipe.run, |
|
cache_examples=CACHE_EXAMPLES, |
|
) |
|
|
|
model_id.change( |
|
fn=app.load_model_info, |
|
inputs=model_id, |
|
outputs=[ |
|
base_model_used_for_training, |
|
prompt_used_for_training, |
|
], |
|
api_name=False, |
|
) |
|
inputs = [ |
|
model_id, |
|
prompt, |
|
video_length, |
|
fps, |
|
seed, |
|
num_steps, |
|
guidance_scale, |
|
] |
|
prompt.submit( |
|
fn=pipe.run, |
|
inputs=inputs, |
|
outputs=result, |
|
api_name=False, |
|
) |
|
run_button.click( |
|
fn=pipe.run, |
|
inputs=inputs, |
|
outputs=result, |
|
api_name="run", |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.queue(max_size=20).launch() |
|
|
