Sebestianmek/open-sora
0
1#!/usr/bin/env python2"""3This script runs a Gradio App for the Open-Sora model.4 5Usage:6 python demo.py <config-path>7"""8 9import argparse10import datetime11import importlib12import os13import subprocess14import sys15from tempfile import NamedTemporaryFile16 17import spaces18import torch19 20import gradio as gr21 22MODEL_TYPES = ["v1.2-stage3"]23WATERMARK_PATH = "./assets/images/watermark/watermark.png"24CONFIG_MAP = {25 "v1.2-stage3": "configs/opensora-v1-2/inference/sample.py",26}27HF_STDIT_MAP = {"v1.2-stage3": "hpcai-tech/OpenSora-STDiT-v3"}28 29 30# ============================31# Prepare Runtime Environment32# ============================33def install_dependencies(enable_optimization=False):34 """35 Install the required dependencies for the demo if they are not already installed.36 """37 38 def _is_package_available(name) -> bool:39 try:40 importlib.import_module(name)41 return True42 except (ImportError, ModuleNotFoundError):43 return False44 45 if enable_optimization:46 # install flash attention47 if not _is_package_available("flash_attn"):48 subprocess.run(49 f"{sys.executable} -m pip install flash-attn --no-build-isolation",50 env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},51 shell=True,52 )53 54 # install apex for fused layernorm55 if not _is_package_available("apex"):56 subprocess.run(57 f'{sys.executable} -m pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" git+https://github.com/NVIDIA/apex.git',58 shell=True,59 )60 61 # install ninja62 if not _is_package_available("ninja"):63 subprocess.run(f"{sys.executable} -m pip install ninja", shell=True)64 65 # install xformers66 if not _is_package_available("xformers"):67 subprocess.run(68 f"{sys.executable} -m pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers",69 shell=True,70 )71 72 73# ============================74# Model-related75# ============================76def read_config(config_path):77 """78 Read the configuration file.79 """80 from mmengine.config import Config81 82 return Config.fromfile(config_path)83 84 85def build_models(model_type, config, enable_optimization=False):86 """87 Build the models for the given model type and configuration.88 """89 # build vae90 from opensora.registry import MODELS, build_module91 92 vae = build_module(config.vae, MODELS).cuda()93 94 # build text encoder95 text_encoder = build_module(config.text_encoder, MODELS) # T5 must be fp3296 text_encoder.t5.model = text_encoder.t5.model.cuda()97 98 # build stdit99 # we load model from HuggingFace directly so that we don't need to100 # handle model download logic in HuggingFace Space101 from opensora.models.stdit.stdit3 import STDiT3102 103 model_kwargs = {k: v for k, v in config.model.items() if k not in ("type", "from_pretrained", "force_huggingface")}104 stdit = STDiT3.from_pretrained(HF_STDIT_MAP[model_type], **model_kwargs)105 stdit = stdit.cuda()106 107 # build scheduler108 from opensora.registry import SCHEDULERS109 110 scheduler = build_module(config.scheduler, SCHEDULERS)111 112 # hack for classifier-free guidance113 text_encoder.y_embedder = stdit.y_embedder114 115 # move modelst to device116 vae = vae.to(torch.bfloat16).eval()117 text_encoder.t5.model = text_encoder.t5.model.eval() # t5 must be in fp32118 stdit = stdit.to(torch.bfloat16).eval()119 120 # clear cuda121 torch.cuda.empty_cache()122 return vae, text_encoder, stdit, scheduler123 124 125def parse_args():126 parser = argparse.ArgumentParser()127 parser.add_argument(128 "--model-type",129 default="v1.2-stage3",130 choices=MODEL_TYPES,131 help=f"The type of model to run for the Gradio App, can only be {MODEL_TYPES}",132 )133 parser.add_argument("--output", default="./outputs", type=str, help="The path to the output folder")134 parser.add_argument("--port", default=None, type=int, help="The port to run the Gradio App on.")135 parser.add_argument("--host", default="0.0.0.0", type=str, help="The host to run the Gradio App on.")136 parser.add_argument("--share", action="store_true", help="Whether to share this gradio demo.")137 parser.add_argument(138 "--enable-optimization",139 action="store_true",140 help="Whether to enable optimization such as flash attention and fused layernorm",141 )142 return parser.parse_args()143 144 145# ============================146# Main Gradio Script147# ============================148# as `run_inference` needs to be wrapped by `spaces.GPU` and the input can only be the prompt text149# so we can't pass the models to `run_inference` as arguments.150# instead, we need to define them globally so that we can access these models inside `run_inference`151 152# read config153args = parse_args()154config = read_config(CONFIG_MAP[args.model_type])155torch.backends.cuda.matmul.allow_tf32 = True156torch.backends.cudnn.allow_tf32 = True157 158# make outputs dir159os.makedirs(args.output, exist_ok=True)160 161# disable torch jit as it can cause failure in gradio SDK162# gradio sdk uses torch with cuda 11.3163torch.jit._state.disable()164 165# set up166install_dependencies(enable_optimization=args.enable_optimization)167 168# import after installation169from opensora.datasets import IMG_FPS, save_sample170from opensora.datasets.aspect import get_image_size, get_num_frames171from opensora.models.text_encoder.t5 import text_preprocessing172from opensora.utils.inference_utils import (173 add_watermark,174 append_generated,175 append_score_to_prompts,176 apply_mask_strategy,177 collect_references_batch,178 dframe_to_frame,179 extract_json_from_prompts,180 extract_prompts_loop,181 get_random_prompt_by_openai,182 has_openai_key,183 merge_prompt,184 prepare_multi_resolution_info,185 refine_prompts_by_openai,186 split_prompt,187)188from opensora.utils.misc import to_torch_dtype189 190# some global variables191dtype = to_torch_dtype(config.dtype)192device = torch.device("cuda")193 194# build model195vae, text_encoder, stdit, scheduler = build_models(196 args.model_type, config, enable_optimization=args.enable_optimization197)198 199 200def run_inference(201 mode,202 prompt_text,203 resolution,204 aspect_ratio,205 length,206 motion_strength,207 aesthetic_score,208 use_motion_strength,209 use_aesthetic_score,210 camera_motion,211 reference_image,212 refine_prompt,213 fps,214 num_loop,215 seed,216 sampling_steps,217 cfg_scale,218):219 if prompt_text is None or prompt_text == "":220 gr.Warning("Your prompt is empty, please enter a valid prompt")221 return None222 223 torch.manual_seed(seed)224 with torch.inference_mode():225 # ======================226 # 1. Preparation arguments227 # ======================228 # parse the inputs229 # frame_interval must be 1 so we ignore it here230 image_size = get_image_size(resolution, aspect_ratio)231 232 # compute generation parameters233 if mode == "Text2Image":234 num_frames = 1235 fps = IMG_FPS236 else:237 num_frames = config.num_frames238 num_frames = get_num_frames(length)239 240 condition_frame_length = int(num_frames / 17 * 5 / 3)241 condition_frame_edit = 0.0242 243 input_size = (num_frames, *image_size)244 latent_size = vae.get_latent_size(input_size)245 multi_resolution = "OpenSora"246 align = 5247 248 # == prepare mask strategy ==249 if mode == "Text2Image":250 mask_strategy = [None]251 elif mode == "Text2Video":252 if reference_image is not None:253 mask_strategy = ["0"]254 else:255 mask_strategy = [None]256 else:257 raise ValueError(f"Invalid mode: {mode}")258 259 # == prepare reference ==260 if mode == "Text2Image":261 refs = [""]262 elif mode == "Text2Video":263 if reference_image is not None:264 # save image to disk265 from PIL import Image266 267 im = Image.fromarray(reference_image)268 temp_file = NamedTemporaryFile(suffix=".png")269 im.save(temp_file.name)270 refs = [temp_file.name]271 else:272 refs = [""]273 else:274 raise ValueError(f"Invalid mode: {mode}")275 276 # == get json from prompts ==277 batch_prompts = [prompt_text]278 batch_prompts, refs, mask_strategy = extract_json_from_prompts(batch_prompts, refs, mask_strategy)279 280 # == get reference for condition ==281 refs = collect_references_batch(refs, vae, image_size)282 283 # == multi-resolution info ==284 model_args = prepare_multi_resolution_info(285 multi_resolution, len(batch_prompts), image_size, num_frames, fps, device, dtype286 )287 288 # == process prompts step by step ==289 # 0. split prompt290 # each element in the list is [prompt_segment_list, loop_idx_list]291 batched_prompt_segment_list = []292 batched_loop_idx_list = []293 for prompt in batch_prompts:294 prompt_segment_list, loop_idx_list = split_prompt(prompt)295 batched_prompt_segment_list.append(prompt_segment_list)296 batched_loop_idx_list.append(loop_idx_list)297 298 # 1. refine prompt by openai299 if refine_prompt:300 # check if openai key is provided301 if not has_openai_key():302 gr.Warning("OpenAI API key is not provided, the prompt will not be enhanced.")303 else:304 for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):305 batched_prompt_segment_list[idx] = refine_prompts_by_openai(prompt_segment_list)306 307 # process scores308 aesthetic_score = aesthetic_score if use_aesthetic_score else None309 motion_strength = motion_strength if use_motion_strength and mode != "Text2Image" else None310 camera_motion = None if camera_motion == "none" or mode == "Text2Image" else camera_motion311 # 2. append score312 for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):313 batched_prompt_segment_list[idx] = append_score_to_prompts(314 prompt_segment_list,315 aes=aesthetic_score,316 flow=motion_strength,317 camera_motion=camera_motion,318 )319 320 # 3. clean prompt with T5321 for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):322 batched_prompt_segment_list[idx] = [text_preprocessing(prompt) for prompt in prompt_segment_list]323 324 # 4. merge to obtain the final prompt325 batch_prompts = []326 for prompt_segment_list, loop_idx_list in zip(batched_prompt_segment_list, batched_loop_idx_list):327 batch_prompts.append(merge_prompt(prompt_segment_list, loop_idx_list))328 329 # =========================330 # Generate image/video331 # =========================332 video_clips = []333 334 for loop_i in range(num_loop):335 # 4.4 sample in hidden space336 batch_prompts_loop = extract_prompts_loop(batch_prompts, loop_i)337 338 # == loop ==339 if loop_i > 0:340 refs, mask_strategy = append_generated(341 vae, video_clips[-1], refs, mask_strategy, loop_i, condition_frame_length, condition_frame_edit342 )343 344 # == sampling ==345 z = torch.randn(len(batch_prompts), vae.out_channels, *latent_size, device=device, dtype=dtype)346 masks = apply_mask_strategy(z, refs, mask_strategy, loop_i, align=align)347 348 # 4.6. diffusion sampling349 # hack to update num_sampling_steps and cfg_scale350 scheduler_kwargs = config.scheduler.copy()351 scheduler_kwargs.pop("type")352 scheduler_kwargs["num_sampling_steps"] = sampling_steps353 scheduler_kwargs["cfg_scale"] = cfg_scale354 355 scheduler.__init__(**scheduler_kwargs)356 samples = scheduler.sample(357 stdit,358 text_encoder,359 z=z,360 prompts=batch_prompts_loop,361 device=device,362 additional_args=model_args,363 progress=True,364 mask=masks,365 )366 samples = vae.decode(samples.to(dtype), num_frames=num_frames)367 video_clips.append(samples)368 369 # =========================370 # Save output371 # =========================372 video_clips = [val[0] for val in video_clips]373 for i in range(1, num_loop):374 video_clips[i] = video_clips[i][:, dframe_to_frame(condition_frame_length) :]375 video = torch.cat(video_clips, dim=1)376 current_datetime = datetime.datetime.now()377 timestamp = current_datetime.timestamp()378 save_path = os.path.join(args.output, f"output_{timestamp}")379 saved_path = save_sample(video, save_path=save_path, fps=24)380 torch.cuda.empty_cache()381 382 # add watermark383 # all watermarked videos should have a _watermarked suffix384 if mode != "Text2Image" and os.path.exists(WATERMARK_PATH):385 watermarked_path = saved_path.replace(".mp4", "_watermarked.mp4")386 success = add_watermark(saved_path, WATERMARK_PATH, watermarked_path)387 if success:388 return watermarked_path389 else:390 return saved_path391 else:392 return saved_path393 394 395@spaces.GPU(duration=200)396def run_image_inference(397 prompt_text,398 resolution,399 aspect_ratio,400 length,401 motion_strength,402 aesthetic_score,403 use_motion_strength,404 use_aesthetic_score,405 camera_motion,406 reference_image,407 refine_prompt,408 fps,409 num_loop,410 seed,411 sampling_steps,412 cfg_scale,413):414 return run_inference(415 "Text2Image",416 prompt_text,417 resolution,418 aspect_ratio,419 length,420 motion_strength,421 aesthetic_score,422 use_motion_strength,423 use_aesthetic_score,424 camera_motion,425 reference_image,426 refine_prompt,427 fps,428 num_loop,429 seed,430 sampling_steps,431 cfg_scale,432 )433 434 435@spaces.GPU(duration=200)436def run_video_inference(437 prompt_text,438 resolution,439 aspect_ratio,440 length,441 motion_strength,442 aesthetic_score,443 use_motion_strength,444 use_aesthetic_score,445 camera_motion,446 reference_image,447 refine_prompt,448 fps,449 num_loop,450 seed,451 sampling_steps,452 cfg_scale,453):454 # if (resolution == "480p" and length == "16s") or \455 # (resolution == "720p" and length in ["8s", "16s"]):456 # gr.Warning("Generation is interrupted as the combination of 480p and 16s will lead to CUDA out of memory")457 # else:458 return run_inference(459 "Text2Video",460 prompt_text,461 resolution,462 aspect_ratio,463 length,464 motion_strength,465 aesthetic_score,466 use_motion_strength,467 use_aesthetic_score,468 camera_motion,469 reference_image,470 refine_prompt,471 fps,472 num_loop,473 seed,474 sampling_steps,475 cfg_scale,476 )477 478 479def generate_random_prompt():480 if "OPENAI_API_KEY" not in os.environ:481 gr.Warning("Your prompt is empty and the OpenAI API key is not provided, please enter a valid prompt")482 return None483 else:484 prompt_text = get_random_prompt_by_openai()485 return prompt_text486 487 488def main():489 # create demo490 with gr.Blocks() as demo:491 with gr.Row():492 with gr.Column():493 gr.HTML(494 """495 <div style='text-align: center;'>496 <p align="center">497 <img src="https://github.com/hpcaitech/Open-Sora/raw/main/assets/readme/icon.png" width="250"/>498 </p>499 <div style="display: flex; gap: 10px; justify-content: center;">500 <a href="https://github.com/hpcaitech/Open-Sora/stargazers"><img src="https://img.shields.io/github/stars/hpcaitech/Open-Sora?style=social"></a>501 <a href="https://hpcaitech.github.io/Open-Sora/"><img src="https://img.shields.io/badge/Gallery-View-orange?logo=&"></a>502 <a href="https://discord.gg/kZakZzrSUT"><img src="https://img.shields.io/badge/Discord-join-blueviolet?logo=discord&"></a>503 <a href="https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-247ipg9fk-KRRYmUl~u2ll2637WRURVA"><img src="https://img.shields.io/badge/Slack-ColossalAI-blueviolet?logo=slack&"></a>504 <a href="https://twitter.com/yangyou1991/status/1769411544083996787?s=61&t=jT0Dsx2d-MS5vS9rNM5e5g"><img src="https://img.shields.io/badge/Twitter-Discuss-blue?logo=twitter&"></a>505 <a href="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png"><img src="https://img.shields.io/badge/微信-小助手加群-green?logo=wechat&"></a>506 <a href="https://hpc-ai.com/blog/open-sora-v1.0"><img src="https://img.shields.io/badge/Open_Sora-Blog-blue"></a>507 </div>508 <h1 style='margin-top: 5px;'>Open-Sora: Democratizing Efficient Video Production for All</h1>509 </div>510 """511 )512 513 with gr.Row():514 with gr.Column():515 prompt_text = gr.Textbox(label="Prompt", placeholder="Describe your video here", lines=4)516 refine_prompt = gr.Checkbox(517 value=has_openai_key(), label="Refine prompt with GPT4o", interactive=has_openai_key()518 )519 random_prompt_btn = gr.Button("Random Prompt By GPT4o", interactive=has_openai_key())520 521 gr.Markdown("## Basic Settings")522 resolution = gr.Radio(523 choices=["144p", "240p", "360p", "480p", "720p"],524 value="480p",525 label="Resolution",526 )527 aspect_ratio = gr.Radio(528 choices=["9:16", "16:9", "3:4", "4:3", "1:1"],529 value="9:16",530 label="Aspect Ratio (H:W)",531 )532 length = gr.Radio(533 choices=["2s", "4s", "8s", "16s"],534 value="2s",535 label="Video Length",536 info="only effective for video generation, 8s may fail as Hugging Face ZeroGPU has the limitation of max 200 seconds inference time.",537 )538 539 with gr.Row():540 seed = gr.Slider(value=1024, minimum=1, maximum=2048, step=1, label="Seed")541 542 sampling_steps = gr.Slider(value=30, minimum=1, maximum=200, step=1, label="Sampling steps")543 cfg_scale = gr.Slider(value=7.0, minimum=0.0, maximum=10.0, step=0.1, label="CFG Scale")544 545 with gr.Row():546 with gr.Column():547 motion_strength = gr.Slider(548 value=5,549 minimum=0,550 maximum=100,551 step=1,552 label="Motion Strength",553 info="only effective for video generation",554 )555 use_motion_strength = gr.Checkbox(value=False, label="Enable")556 557 with gr.Column():558 aesthetic_score = gr.Slider(559 value=6.5,560 minimum=4,561 maximum=7,562 step=0.1,563 label="Aesthetic",564 info="effective for text & video generation",565 )566 use_aesthetic_score = gr.Checkbox(value=True, label="Enable")567 568 camera_motion = gr.Radio(569 value="none",570 label="Camera Motion",571 choices=["none", "pan right", "pan left", "tilt up", "tilt down", "zoom in", "zoom out", "static"],572 interactive=True,573 )574 575 gr.Markdown("## Advanced Settings")576 with gr.Row():577 fps = gr.Slider(578 value=24,579 minimum=1,580 maximum=60,581 step=1,582 label="FPS",583 info="This is the frames per seconds for video generation, keep it to 24 if you are not sure",584 )585 num_loop = gr.Slider(586 value=1,587 minimum=1,588 maximum=20,589 step=1,590 label="Number of Loops",591 info="This will change the length of the generated video, keep it to 1 if you are not sure",592 )593 594 gr.Markdown("## Reference Image")595 reference_image = gr.Image(label="Image (optional)", show_download_button=True)596 597 with gr.Column():598 output_video = gr.Video(label="Output Video", height="100%")599 600 with gr.Row():601 image_gen_button = gr.Button("Generate image")602 video_gen_button = gr.Button("Generate video")603 604 image_gen_button.click(605 fn=run_image_inference,606 inputs=[607 prompt_text,608 resolution,609 aspect_ratio,610 length,611 motion_strength,612 aesthetic_score,613 use_motion_strength,614 use_aesthetic_score,615 camera_motion,616 reference_image,617 refine_prompt,618 fps,619 num_loop,620 seed,621 sampling_steps,622 cfg_scale,623 ],624 outputs=reference_image,625 )626 video_gen_button.click(627 fn=run_video_inference,628 inputs=[629 prompt_text,630 resolution,631 aspect_ratio,632 length,633 motion_strength,634 aesthetic_score,635 use_motion_strength,636 use_aesthetic_score,637 camera_motion,638 reference_image,639 refine_prompt,640 fps,641 num_loop,642 seed,643 sampling_steps,644 cfg_scale,645 ],646 outputs=output_video,647 )648 random_prompt_btn.click(fn=generate_random_prompt, outputs=prompt_text)649 650 # launch651 demo.queue(max_size=5, default_concurrency_limit=1)652 demo.launch(server_port=args.port, server_name=args.host, share=args.share, max_threads=1)653 654 655if __name__ == "__main__":656 main()657