Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
from diffusers import StableDiffusionPipeline, DDIMScheduler | |
from utils import video_to_frames, add_dict_to_yaml_file, save_video, seed_everything | |
# from diffusers.utils import export_to_video | |
from tokenflow_pnp import TokenFlow | |
from preprocess_utils import * | |
from tokenflow_utils import * | |
# load sd model | |
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
if torch.cuda.is_available(): | |
device = "cuda" | |
elif torch.backends.mps.is_available(): | |
device = "mps" | |
else: | |
device = "cpu" | |
model_id = "stabilityai/stable-diffusion-2-1-base" | |
to = torch.float16 if device == 'cuda' else torch.float32 | |
# components for the Preprocessor | |
scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler") | |
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", revision="fp16", | |
torch_dtype=to).to(device) | |
tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") | |
text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision="fp16", | |
torch_dtype=to).to(device) | |
unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", revision="fp16", | |
torch_dtype=to).to(device) | |
# pipe for TokenFlow | |
tokenflow_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=to).to(device) | |
if device == "cuda": | |
tokenflow_pipe.enable_xformers_memory_efficient_attention() | |
def randomize_seed_fn(): | |
seed = random.randint(0, np.iinfo(np.int32).max) | |
return seed | |
def reset_do_inversion(): | |
return True | |
def get_example(): | |
case = [ | |
[ | |
'examples/wolf.mp4', | |
], | |
[ | |
'examples/woman-running.mp4', | |
], | |
[ | |
'examples/cutting_bread.mp4', | |
], | |
[ | |
'examples/running_dog.mp4', | |
] | |
] | |
return case | |
def prep(config): | |
# timesteps to save | |
if config["sd_version"] == '2.1': | |
model_key = "stabilityai/stable-diffusion-2-1-base" | |
elif config["sd_version"] == '2.0': | |
model_key = "stabilityai/stable-diffusion-2-base" | |
elif config["sd_version"] == '1.5' or config["sd_version"] == 'ControlNet': | |
model_key = "runwayml/stable-diffusion-v1-5" | |
elif config["sd_version"] == 'depth': | |
model_key = "stabilityai/stable-diffusion-2-depth" | |
toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler") | |
toy_scheduler.set_timesteps(config["save_steps"]) | |
print("config[save_steps]", config["save_steps"]) | |
timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=config["save_steps"], | |
strength=1.0, | |
device=device) | |
print("YOOOO timesteps to save", timesteps_to_save) | |
# seed_everything(config["seed"]) | |
if not config["frames"]: # original non demo setting | |
save_path = os.path.join(config["save_dir"], | |
f'sd_{config["sd_version"]}', | |
Path(config["data_path"]).stem, | |
f'steps_{config["steps"]}', | |
f'nframes_{config["n_frames"]}') | |
os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True) | |
add_dict_to_yaml_file(os.path.join(config["save_dir"], 'inversion_prompts.yaml'), Path(config["data_path"]).stem, config["inversion_prompt"]) | |
# save inversion prompt in a txt file | |
with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f: | |
f.write(config["inversion_prompt"]) | |
else: | |
save_path = None | |
model = Preprocess(device, config, | |
vae=vae, | |
text_encoder=text_encoder, | |
scheduler=scheduler, | |
tokenizer=tokenizer, | |
unet=unet) | |
print(type(model.config["batch_size"])) | |
frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents( | |
num_steps=model.config["steps"], | |
save_path=save_path, | |
batch_size=model.config["batch_size"], | |
timesteps_to_save=timesteps_to_save, | |
inversion_prompt=model.config["inversion_prompt"], | |
) | |
return frames, latents, total_inverted_latents, rgb_reconstruction | |
def preprocess_and_invert(input_video, | |
frames, | |
latents, | |
inverted_latents, | |
seed, | |
randomize_seed, | |
do_inversion, | |
# save_dir: str = "latents", | |
steps, | |
n_timesteps = 50, | |
batch_size: int = 8, | |
n_frames: int = 40, | |
inversion_prompt:str = '', | |
): | |
sd_version = "2.1" | |
height = 512 | |
weidth: int = 512 | |
print("n timesteps", n_timesteps) | |
if do_inversion or randomize_seed: | |
preprocess_config = {} | |
preprocess_config['H'] = height | |
preprocess_config['W'] = weidth | |
preprocess_config['save_dir'] = 'latents' | |
preprocess_config['sd_version'] = sd_version | |
preprocess_config['steps'] = steps | |
preprocess_config['batch_size'] = batch_size | |
preprocess_config['save_steps'] = int(n_timesteps) | |
preprocess_config['n_frames'] = n_frames | |
preprocess_config['seed'] = seed | |
preprocess_config['inversion_prompt'] = inversion_prompt | |
preprocess_config['frames'] = video_to_frames(input_video) | |
preprocess_config['data_path'] = input_video.split(".")[0] | |
if randomize_seed: | |
seed = randomize_seed_fn() | |
seed_everything(seed) | |
frames, latents, total_inverted_latents, rgb_reconstruction = prep(preprocess_config) | |
print(total_inverted_latents.keys()) | |
print(len(total_inverted_latents.keys())) | |
frames = gr.State(value=frames) | |
latents = gr.State(value=latents) | |
inverted_latents = gr.State(value=total_inverted_latents) | |
do_inversion = False | |
return frames, latents, inverted_latents, do_inversion | |
def edit_with_pnp(input_video, | |
frames, | |
latents, | |
inverted_latents, | |
seed, | |
randomize_seed, | |
do_inversion, | |
steps, | |
prompt: str = "a marble sculpture of a woman running, Venus de Milo", | |
# negative_prompt: str = "ugly, blurry, low res, unrealistic, unaesthetic", | |
pnp_attn_t: float = 0.5, | |
pnp_f_t: float = 0.8, | |
batch_size: int = 8, #needs to be the same as for preprocess | |
n_frames: int = 40,#needs to be the same as for preprocess | |
n_timesteps: int = 50, | |
gudiance_scale: float = 7.5, | |
inversion_prompt: str = "", #needs to be the same as for preprocess | |
n_fps: int = 10, | |
progress=gr.Progress(track_tqdm=True) | |
): | |
config = {} | |
config["sd_version"] = "2.1" | |
config["device"] = device | |
config["n_timesteps"] = int(n_timesteps) | |
config["n_frames"] = n_frames | |
config["batch_size"] = batch_size | |
config["guidance_scale"] = gudiance_scale | |
config["prompt"] = prompt | |
config["negative_prompt"] = "ugly, blurry, low res, unrealistic, unaesthetic", | |
config["pnp_attn_t"] = pnp_attn_t | |
config["pnp_f_t"] = pnp_f_t | |
config["pnp_inversion_prompt"] = inversion_prompt | |
if do_inversion: | |
frames, latents, inverted_latents, do_inversion = preprocess_and_invert( | |
input_video, | |
frames, | |
latents, | |
inverted_latents, | |
seed, | |
randomize_seed, | |
do_inversion, | |
steps, | |
n_timesteps, | |
batch_size, | |
n_frames, | |
inversion_prompt) | |
do_inversion = False | |
if randomize_seed: | |
seed = randomize_seed_fn() | |
seed_everything(seed) | |
editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value) | |
edited_frames = editor.edit_video() | |
save_video(edited_frames, 'tokenflow_PnP_fps_30.mp4', fps=n_fps) | |
# path = export_to_video(edited_frames) | |
return 'tokenflow_PnP_fps_30.mp4', frames, latents, inverted_latents, do_inversion | |
######## | |
# demo # | |
######## | |
intro = """ | |
<div style="text-align:center"> | |
<h1 style="font-weight: 1400; text-align: center; margin-bottom: 7px;"> | |
TokenFlow - <small>Temporally consistent video editing</small> | |
</h1> | |
<span>[<a target="_blank" href="https://diffusion-tokenflow.github.io">Project page</a>], [<a target="_blank" href="https://github.com/omerbt/TokenFlow">GitHub</a>], [<a target="_blank" href="https://huggingface.co/papers/2307.10373">Paper</a>]</span> | |
<div style="display:flex; justify-content: center;margin-top: 0.5em">Each edit takes ~5 min <a href="https://huggingface.co/weizmannscience/tokenflow?duplicate=true" target="_blank"> | |
<img style="margin-top: 0em; margin-bottom: 0em; margin-left: 0.5em" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a></div> | |
</div> | |
""" | |
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_video = gr.Video(label="Input Video", interactive=True, elem_id="input_video") | |
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) | |
with gr.Row(): | |
prompt = gr.Textbox( | |
label="Describe your edited video", | |
max_lines=1, value="" | |
) | |
# 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'): | |
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) | |
gudiance_scale = gr.Slider(label='Guidance Scale', minimum=1, maximum=30, | |
value=7.5, step=0.5, interactive=True) | |
steps = gr.Slider(label='Inversion steps', minimum=10, maximum=500, | |
value=200, 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=2, maximum=200, | |
value=24, step=1, interactive=True) | |
n_timesteps = gr.Slider(label='Diffusion steps', minimum=25, maximum=100, | |
value=25, step=25, interactive=True) | |
n_fps = gr.Slider(label='Frames per second', minimum=1, maximum=60, | |
value=10, step=1, interactive=True) | |
with gr.TabItem('Plug-and-Play Parameters'): | |
with gr.Column(min_width=100): | |
pnp_attn_t = gr.Slider(label='pnp attention threshold', minimum=0, maximum=1, | |
value=0.5, step=0.5, interactive=True) | |
pnp_f_t = gr.Slider(label='pnp feature threshold', minimum=0, maximum=1, | |
value=0.8, step=0.05, interactive=True) | |
input_video.change( | |
fn = reset_do_inversion, | |
outputs = [do_inversion], | |
queue = False) | |
inversion_prompt.change( | |
fn = reset_do_inversion, | |
outputs = [do_inversion], | |
queue = False) | |
randomize_seed.change( | |
fn = reset_do_inversion, | |
outputs = [do_inversion], | |
queue = False) | |
seed.change( | |
fn = reset_do_inversion, | |
outputs = [do_inversion], | |
queue = False) | |
input_video.upload( | |
fn = reset_do_inversion, | |
outputs = [do_inversion], | |
queue = False).then(fn = preprocess_and_invert, | |
inputs = [input_video, | |
frames, | |
latents, | |
inverted_latents, | |
seed, | |
randomize_seed, | |
do_inversion, | |
steps, | |
n_timesteps, | |
batch_size, | |
n_frames, | |
inversion_prompt | |
], | |
outputs = [frames, | |
latents, | |
inverted_latents, | |
do_inversion | |
]) | |
run_button.click(fn = edit_with_pnp, | |
inputs = [input_video, | |
frames, | |
latents, | |
inverted_latents, | |
seed, | |
randomize_seed, | |
do_inversion, | |
steps, | |
prompt, | |
pnp_attn_t, | |
pnp_f_t, | |
batch_size, | |
n_frames, | |
n_timesteps, | |
gudiance_scale, | |
inversion_prompt, | |
n_fps ], | |
outputs = [output_video, frames, latents, inverted_latents, do_inversion] | |
) | |
gr.Examples( | |
examples=get_example(), | |
label='Examples', | |
inputs=[input_video], | |
outputs=[output_video] | |
) | |
demo.queue() | |
demo.launch() | |