Spaces:
Runtime error
Runtime error
File size: 5,747 Bytes
68aadbb 87e3aec 68aadbb 87e3aec 68aadbb 87e3aec 68aadbb 87e3aec 68aadbb 87e3aec 68aadbb 87e3aec 68aadbb 87e3aec d7d00c9 68aadbb d7d00c9 fa2a666 d7d00c9 dd28623 d7d00c9 68aadbb d7d00c9 68aadbb d7d00c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler
from utils import *
# load sd model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = "stabilityai/stable-diffusion-2-1-base"
inv_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
inv_pipe.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
def randomize_seed_fn():
seed = random.randint(0, np.iinfo(np.int32).max)
return seed
def preprocess_and_invert(video,
frames,
latents,
inverted_latents,
seed,
randomize_seed,
do_inversion,
height:int = 512,
weidth: int = 512,
# save_dir: str = "latents",
steps: int = 500,
batch_size: int = 8,
# save_steps: int = 50,
n_frames: int = 40,
inversion_prompt:str = ''
):
if do_inversion or randomize_seed:
# save_video_frames(data_path, img_size=(height, weidth))
frames = video_to_frames(video, img_size=(height, weidth))
# data_path = os.path.join('data', Path(video_path).stem)
toy_scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
toy_scheduler.set_timesteps(save_steps)
timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=save_steps,
strength=1.0,
device=device)
if randomize_seed:
seed = randomize_seed_fn()
seed_everything(seed)
frames, latents = get_data(inv_pipe, frames, n_frames)
inverted_latents = extract_latents(inv_pipe, num_steps = steps,
latent_frames = latents,
batch_size = batch_size,
timesteps_to_save = timesteps_to_save,
inversion_prompt = inversion_prompt,)
frames = gr.State(value=frames)
latents = gr.State(value=latents)
inverted_latents = gr.State(value=inverted_latents)
do_inversion = False
return frames, latents, inverted_latents, do_inversion
########
# demo #
########
intro = """
<h1 style="font-weight: 1400; text-align: center; margin-bottom: 7px;">
TokenFlow
</h1>
"""
with gr.Blocks(css="style.css") as demo:
gr.HTML(intro)
frames = gr.State()
inverted_latents = gr.State()
latents = gr.State()
do_inversion = gr.State(value=True)
with gr.Row():
input_vid = gr.Video(label="Input Video", interactive=True, elem_id="input_video")
output_vid = gr.Video(label="Edited Video", interactive=False, elem_id="output_video")
input_vid.style(height=365, width=365)
output_vid.style(height=365, width=365)
# with gr.Group(visible=False) as share_btn_container:
# with gr.Group(elem_id="share-btn-container"):
# community_icon = gr.HTML(community_icon_html, visible=True)
# loading_icon = gr.HTML(loading_icon_html, visible=False)
# share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
# with gr.Row():
# inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
with gr.Row():
run_button = gr.Button("Edit your video!", visible=True)
with gr.Accordion("Advanced Options", open=False):
with gr.Tabs() as tabs:
with gr.TabItem('General options', id=2):
with gr.Row():
with gr.Column(min_width=100):
seed = gr.Number(value=0, precision=0, label="Seed", interactive=True)
randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
steps = gr.Slider(label='Inversion steps', minimum=100, maximum=500,
value=500, step=1, interactive=True)
with gr.Column(min_width=100):
inversion_prompt = gr.Textbox(lines=1, label="Inversion prompt", interactive=True, placeholder="")
batch_size = gr.Slider(label='Batch size', minimum=1, maximum=10,
value=8, step=1, interactive=True)
n_frames = gr.Slider(label='Num frames', minimum=20, maximum=200,
value=40, step=1, interactive=True)
input_vid.change(
fn = reset_do_inversion,
outputs = [do_inversion],
queue = False)
input_vid.upload(
fn = reset_do_inversion,
outputs = [do_inversion],
queue = False)
).then(fn = preprocess_and_invert,
inputs = [input_vid,
frames,
latents,
inverted_latents,
seed,
randomize_seed,
do_inversion,
steps,
batch_size,
n_frames,
inversion_prompt
],
outputs = [frames,
latents,
inverted_latents,
do_inversion
])
demo.queue()
demo.launch() |