update

app.py

@@ -11,147 +11,25 @@ import importlib
 import os
 import subprocess
 import sys
-import re
-import json
-import math
-
 import spaces
 import torch
 
 import gradio as gr
+from tempfile import NamedTemporaryFile
+import datetime
+
 
 
-MODEL_TYPES = ["v1.1"]
+MODEL_TYPES = ["v1.2-stage3"]
 CONFIG_MAP = {
-    "v1.1-stage2": "configs/opensora-v1-1/inference/sample-ref.py",
-    "v1.1-stage3": "configs/opensora-v1-1/inference/sample-ref.py",
+    "v1.2-stage3": "configs/opensora-v1-2/inference/sample.py",
 }
 HF_STDIT_MAP = {
-    "v1.1-stage2": "hpcai-tech/OpenSora-STDiT-v2-stage2",
-    "v1.1-stage3": "hpcai-tech/OpenSora-STDiT-v2-stage3",
-}
-RESOLUTION_MAP = {
-    "144p": (144, 256),
-    "240p": (240, 426),
-    "360p": (360, 480),
-    "480p": (480, 858),
-    "720p": (720, 1280),
-    "1080p": (1080, 1920)
+    "v1.2-stage3": "hpcai-tech/OpenSora-STDiT-v3"
 }
 
-
-# ============================
-# Utils
 # ============================
-def collect_references_batch(reference_paths, vae, image_size):
-    from opensora.datasets.utils import read_from_path
-
-    refs_x = []
-    for reference_path in reference_paths:
-        if reference_path is None:
-            refs_x.append([])
-            continue
-        ref_path = reference_path.split(";")
-        ref = []
-        for r_path in ref_path:
-            r = read_from_path(r_path, image_size, transform_name="resize_crop")
-            r_x = vae.encode(r.unsqueeze(0).to(vae.device, vae.dtype))
-            r_x = r_x.squeeze(0)
-            ref.append(r_x)
-        refs_x.append(ref)
-    # refs_x: [batch, ref_num, C, T, H, W]
-    return refs_x
-
-
-def process_mask_strategy(mask_strategy):
-    mask_batch = []
-    mask_strategy = mask_strategy.split(";")
-    for mask in mask_strategy:
-        mask_group = mask.split(",")
-        assert len(mask_group) >= 1 and len(mask_group) <= 6, f"Invalid mask strategy: {mask}"
-        if len(mask_group) == 1:
-            mask_group.extend(["0", "0", "0", "1", "0"])
-        elif len(mask_group) == 2:
-            mask_group.extend(["0", "0", "1", "0"])
-        elif len(mask_group) == 3:
-            mask_group.extend(["0", "1", "0"])
-        elif len(mask_group) == 4:
-            mask_group.extend(["1", "0"])
-        elif len(mask_group) == 5:
-            mask_group.append("0")
-        mask_batch.append(mask_group)
-    return mask_batch
-
-
-def apply_mask_strategy(z, refs_x, mask_strategys, loop_i):
-    masks = []
-    for i, mask_strategy in enumerate(mask_strategys):
-        mask = torch.ones(z.shape[2], dtype=torch.float, device=z.device)
-        if mask_strategy is None:
-            masks.append(mask)
-            continue
-        mask_strategy = process_mask_strategy(mask_strategy)
-        for mst in mask_strategy:
-            loop_id, m_id, m_ref_start, m_target_start, m_length, edit_ratio = mst
-            loop_id = int(loop_id)
-            if loop_id != loop_i:
-                continue
-            m_id = int(m_id)
-            m_ref_start = int(m_ref_start)
-            m_length = int(m_length)
-            m_target_start = int(m_target_start)
-            edit_ratio = float(edit_ratio)
-            ref = refs_x[i][m_id]  # [C, T, H, W]
-            if m_ref_start < 0:
-                m_ref_start = ref.shape[1] + m_ref_start
-            if m_target_start < 0:
-                # z: [B, C, T, H, W]
-                m_target_start = z.shape[2] + m_target_start
-            z[i, :, m_target_start : m_target_start + m_length] = ref[:, m_ref_start : m_ref_start + m_length]
-            mask[m_target_start : m_target_start + m_length] = edit_ratio
-        masks.append(mask)
-    masks = torch.stack(masks)
-    return masks
-
-
-def process_prompts(prompts, num_loop):
-    from opensora.models.text_encoder.t5 import text_preprocessing
-
-    ret_prompts = []
-    for prompt in prompts:
-        if prompt.startswith("|0|"):
-            prompt_list = prompt.split("|")[1:]
-            text_list = []
-            for i in range(0, len(prompt_list), 2):
-                start_loop = int(prompt_list[i])
-                text = prompt_list[i + 1]
-                text = text_preprocessing(text)
-                end_loop = int(prompt_list[i + 2]) if i + 2 < len(prompt_list) else num_loop
-                text_list.extend([text] * (end_loop - start_loop))
-            assert len(text_list) == num_loop, f"Prompt loop mismatch: {len(text_list)} != {num_loop}"
-            ret_prompts.append(text_list)
-        else:
-            prompt = text_preprocessing(prompt)
-            ret_prompts.append([prompt] * num_loop)
-    return ret_prompts
-
-
-def extract_json_from_prompts(prompts):
-    additional_infos = []
-    ret_prompts = []
-    for prompt in prompts:
-        parts = re.split(r"(?=[{\[])", prompt)
-        assert len(parts) <= 2, f"Invalid prompt: {prompt}"
-        ret_prompts.append(parts[0])
-        if len(parts) == 1:
-            additional_infos.append({})
-        else:
-            additional_infos.append(json.loads(parts[1]))
-    return ret_prompts, additional_infos
-
-
-# ============================
-# Runtime Environment
+# Prepare Runtime Environment
 # ============================
 def install_dependencies(enable_optimization=False):
     """
@@ -223,13 +101,9 @@ def build_models(model_type, config, enable_optimization=False):
     # build stdit
     # we load model from HuggingFace directly so that we don't need to
     # handle model download logic in HuggingFace Space
-    from transformers import AutoModel
-
-    stdit = AutoModel.from_pretrained(
-        HF_STDIT_MAP[model_type],
-        enable_flash_attn=enable_optimization,
-        trust_remote_code=True,
-    ).cuda()
+    from opensora.models.stdit.stdit3 import STDiT3
+    stdit = STDiT3.from_pretrained(HF_STDIT_MAP[model_type])
+    stdit = stdit.cuda()
 
     # build scheduler
     from opensora.registry import SCHEDULERS
@@ -253,13 +127,13 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--model-type",
-        default="v1.1-stage3",
+        default="v1.2-stage3",
         choices=MODEL_TYPES,
         help=f"The type of model to run for the Gradio App, can only be {MODEL_TYPES}",
     )
     parser.add_argument("--output", default="./outputs", type=str, help="The path to the output folder")
     parser.add_argument("--port", default=None, type=int, help="The port to run the Gradio App on.")
-    parser.add_argument("--host", default=None, type=str, help="The host to run the Gradio App on.")
+    parser.add_argument("--host", default="0.0.0.0", type=str, help="The host to run the Gradio App on.")
     parser.add_argument("--share", action="store_true", help="Whether to share this gradio demo.")
     parser.add_argument(
         "--enable-optimization",
@@ -279,6 +153,8 @@ def parse_args():
 # read config
 args = parse_args()
 config = read_config(CONFIG_MAP[args.model_type])
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
 
 # make outputs dir
 os.makedirs(args.output, exist_ok=True)
@@ -293,6 +169,24 @@ install_dependencies(enable_optimization=args.enable_optimization)
 # import after installation
 from opensora.datasets import IMG_FPS, save_sample
 from opensora.utils.misc import to_torch_dtype
+from opensora.utils.inference_utils import (
+    append_generated,
+    apply_mask_strategy,
+    collect_references_batch,
+    extract_json_from_prompts,
+    extract_prompts_loop,
+    prepare_multi_resolution_info,
+    dframe_to_frame,
+    append_score_to_prompts,
+    has_openai_key,
+    refine_prompts_by_openai,
+    add_watermark,
+    get_random_prompt_by_openai,
+    split_prompt,
+    merge_prompt
+)
+from opensora.models.text_encoder.t5 import text_preprocessing
+from opensora.datasets.aspect import get_image_size, get_num_frames
 
 # some global variables
 dtype = to_torch_dtype(config.dtype)
@@ -302,117 +196,278 @@ device = torch.device("cuda")
 vae, text_encoder, stdit, scheduler = build_models(args.model_type, config, enable_optimization=args.enable_optimization)
 
 
-@spaces.GPU(duration=200)
-def run_inference(mode, prompt_text, resolution, length, reference_image):
+def run_inference(mode, prompt_text, resolution, aspect_ratio, length, motion_strength, aesthetic_score, use_motion_strength, use_aesthetic_score, camera_motion, reference_image, refine_prompt, fps, num_loop, seed, sampling_steps, cfg_scale):
+    if prompt_text is None or prompt_text == "":
+        gr.Warning("Your prompt is empty, please enter a valid prompt")
+        return None
+    
+    torch.manual_seed(seed)
     with torch.inference_mode():
         # ======================
-        # 1. Preparation
+        # 1. Preparation arguments
         # ======================
         # parse the inputs
-        resolution = RESOLUTION_MAP[resolution]
+        # frame_interval must be 1 so  we ignore it here
+        image_size = get_image_size(resolution, aspect_ratio)
+        
+        # compute generation parameters
+        if mode == "Text2Image":
+            num_frames = 1
+            fps = IMG_FPS
+        else:
+            num_frames = config.num_frames
+            num_frames = get_num_frames(length)
+            
+        condition_frame_length = int(num_frames / 17 * 5 / 3)
+        condition_frame_edit = 0.0
         
-        # compute number of loops
-        num_seconds = int(length.rstrip('s'))
-        total_number_of_frames = num_seconds * config.fps / config.frame_interval
-        num_loop = math.ceil(total_number_of_frames / config.num_frames)
-
-        # prepare model args
-        model_args = dict()
-        height = torch.tensor([resolution[0]], device=device, dtype=dtype)
-        width = torch.tensor([resolution[1]], device=device, dtype=dtype)
-        num_frames = torch.tensor([config.num_frames], device=device, dtype=dtype)
-        ar = torch.tensor([resolution[0] / resolution[1]], device=device, dtype=dtype)
-        if config.num_frames == 1:
-            config.fps = IMG_FPS
-        fps = torch.tensor([config.fps], device=device, dtype=dtype)
-        model_args["height"] = height
-        model_args["width"] = width
-        model_args["num_frames"] = num_frames
-        model_args["ar"] = ar
-        model_args["fps"] = fps
-
-        # compute latent size
-        input_size = (config.num_frames, *resolution)
+        input_size = (num_frames, *image_size)
         latent_size = vae.get_latent_size(input_size)
-
-        # process prompt
-        prompt_raw = [prompt_text]
-        prompt_raw, _ = extract_json_from_prompts(prompt_raw)
-        prompt_loops = process_prompts(prompt_raw, num_loop)
-        video_clips = []
-
-        # prepare mask strategy
-        if mode == "Text2Video":
+        multi_resolution = "OpenSora"
+        align = 5
+        
+        # == prepare mask strategy ==
+        if mode == "Text2Image":
             mask_strategy = [None]
-        elif mode == "Image2Video":
-            mask_strategy = ['0']
+        elif mode == "Text2Video":
+            if reference_image is not None:
+                mask_strategy = ['0']
+            else:
+                mask_strategy = [None]
         else:
             raise ValueError(f"Invalid mode: {mode}")
-
-        # =========================
-        # 2. Load reference images
-        # =========================
-        if mode == "Text2Video":
-            refs_x = collect_references_batch([None], vae, resolution)
-        elif mode == "Image2Video":
-            # save image to disk
-            from PIL import Image
-            im = Image.fromarray(reference_image)
-            im.save("test.jpg")
-            refs_x = collect_references_batch(["test.jpg"], vae, resolution)
+        
+        # == prepare reference ==
+        if mode == "Text2Image":
+            refs = [""]
+        elif mode == "Text2Video":
+            if reference_image is not None:
+                # save image to disk
+                from PIL import Image
+                im = Image.fromarray(reference_image)
+                temp_file = NamedTemporaryFile(suffix=".png")
+                im.save(temp_file.name)
+                refs = [temp_file.name]
+            else:
+                refs = [""]
         else:
             raise ValueError(f"Invalid mode: {mode}")
+        
+        # == get json from prompts ==
+        batch_prompts = [prompt_text]
+        batch_prompts, refs, mask_strategy = extract_json_from_prompts(batch_prompts, refs, mask_strategy)
 
-        # 4.3. long video generation
-        for loop_i in range(num_loop):
-            # 4.4 sample in hidden space
-            batch_prompts = [prompt[loop_i] for prompt in prompt_loops]
-            z = torch.randn(len(batch_prompts), vae.out_channels, *latent_size, device=device, dtype=dtype)
+        # == get reference for condition ==
+        refs = collect_references_batch(refs, vae, image_size)
 
-            # 4.5. apply mask strategy
-            masks = None
+        # == multi-resolution info ==
+        model_args = prepare_multi_resolution_info(
+            multi_resolution, len(batch_prompts), image_size, num_frames, fps, device, dtype
+        )
+        
+        # == process prompts step by step == 
+        # 0. split prompt
+        # each element in the list is [prompt_segment_list, loop_idx_list]
+        batched_prompt_segment_list = []
+        batched_loop_idx_list = []
+        for prompt in batch_prompts:
+            prompt_segment_list, loop_idx_list = split_prompt(prompt) 
+            batched_prompt_segment_list.append(prompt_segment_list)
+            batched_loop_idx_list.append(loop_idx_list)
+        
+        # 1. refine prompt by openai
+        if refine_prompt:
+            # check if openai key is provided
+            if not has_openai_key():
+                gr.Warning("OpenAI API key is not provided, the prompt will not be enhanced.")
+            else:
+                for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
+                    batched_prompt_segment_list[idx] = refine_prompts_by_openai(prompt_segment_list)
+        
+        # process scores
+        aesthetic_score = aesthetic_score if use_aesthetic_score else None
+        motion_strength = motion_strength if use_motion_strength and mode != "Text2Image" else None
+        camera_motion = None if camera_motion == "none" or mode == "Text2Image" else camera_motion
+        # 2. append score
+        for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
+            batched_prompt_segment_list[idx] = append_score_to_prompts(
+                prompt_segment_list,
+                aes=aesthetic_score,
+                flow=motion_strength,
+                camera_motion=camera_motion,
+            )
 
-            # if cfg.reference_path is not None:
-            if loop_i > 0:
-                ref_x = vae.encode(video_clips[-1])
-                for j, refs in enumerate(refs_x):
-                    if refs is None:
-                        refs_x[j] = [ref_x[j]]
-                    else:
-                        refs.append(ref_x[j])
-                    if mask_strategy[j] is None:
-                        mask_strategy[j] = ""
-                    else:
-                        mask_strategy[j] += ";"
-                    mask_strategy[
-                        j
-                    ] += f"{loop_i},{len(refs)-1},-{config.condition_frame_length},0,{config.condition_frame_length}"
-
-            masks = apply_mask_strategy(z, refs_x, mask_strategy, loop_i)
+        # 3. clean prompt with T5
+        for idx, prompt_segment_list in enumerate(batched_prompt_segment_list):
+            batched_prompt_segment_list[idx] = [text_preprocessing(prompt) for prompt in prompt_segment_list]
+        
+        # 4. merge to obtain the final prompt
+        batch_prompts = []
+        for prompt_segment_list, loop_idx_list in zip(batched_prompt_segment_list, batched_loop_idx_list):
+            batch_prompts.append(merge_prompt(prompt_segment_list, loop_idx_list))
+        
 
+        # =========================
+        # Generate image/video
+        # =========================
+        video_clips = []
+        
+        for loop_i in range(num_loop):
+            # 4.4 sample in hidden space
+            batch_prompts_loop = extract_prompts_loop(batch_prompts, loop_i)
+            
+            # == loop ==
+            if loop_i > 0:
+                refs, mask_strategy = append_generated(
+                    vae, 
+                    video_clips[-1],
+                    refs,
+                    mask_strategy,
+                    loop_i,
+                    condition_frame_length,
+                    condition_frame_edit
+                    )
+            
+            # == sampling ==
+            z = torch.randn(len(batch_prompts), vae.out_channels, *latent_size, device=device, dtype=dtype)
+            masks = apply_mask_strategy(z, refs, mask_strategy, loop_i, align=align)
+            
             # 4.6. diffusion sampling
+            # hack to update num_sampling_steps and cfg_scale
+            scheduler_kwargs = config.scheduler.copy()
+            scheduler_kwargs.pop('type')
+            scheduler_kwargs['num_sampling_steps'] = sampling_steps
+            scheduler_kwargs['cfg_scale'] = cfg_scale
+
+            scheduler.__init__(
+                **scheduler_kwargs
+            )
             samples = scheduler.sample(
                 stdit,
                 text_encoder,
                 z=z,
-                prompts=batch_prompts,
+                prompts=batch_prompts_loop,
                 device=device,
                 additional_args=model_args,
-                mask=masks,  # scheduler must support mask
+                progress=True,
+                mask=masks,
             )
-            samples = vae.decode(samples.to(dtype))
+            samples = vae.decode(samples.to(dtype), num_frames=num_frames)
             video_clips.append(samples)
-
-            # 4.7. save video
-            if loop_i == num_loop - 1:
-                video_clips_list = [
-                    video_clips[0][0]] + [video_clips[i][0][:, config.condition_frame_length :] 
-                    for i in range(1, num_loop)
-                ]
-                video = torch.cat(video_clips_list, dim=1)
-                save_path = f"{args.output}/sample"
-                saved_path = save_sample(video, fps=config.fps // config.frame_interval, save_path=save_path, force_video=True)
+            
+        # =========================
+        # Save output
+        # =========================
+        video_clips = [val[0] for val in video_clips]
+        for i in range(1, num_loop):
+            video_clips[i] = video_clips[i][:, dframe_to_frame(condition_frame_length) :]
+        video = torch.cat(video_clips, dim=1)
+        current_datetime = datetime.datetime.now()
+        timestamp = current_datetime.timestamp()
+        save_path = os.path.join(args.output, f"output_{timestamp}")
+        saved_path = save_sample(video, save_path=save_path, fps=24)
+        torch.cuda.empty_cache()
+        
+        # add watermark
+        # all watermarked videos should have a _watermarked suffix
+        if mode != "Text2Image":
+            watermarked_path = saved_path.replace(".mp4", "_watermarked.mp4")
+            success = add_watermark(saved_path, "./assets/images/watermark/watermark.png", watermarked_path)
+            if success:
+                return watermarked_path
+            else:
                 return saved_path
+        else:
+            return saved_path
+        
+        
+@spaces.GPU(duration=200)
+def run_image_inference(
+    prompt_text, 
+    resolution, 
+    aspect_ratio, 
+    length, 
+    motion_strength, 
+    aesthetic_score, 
+    use_motion_strength, 
+    use_aesthetic_score,
+    camera_motion,
+    reference_image,
+    refine_prompt,
+    fps,
+    num_loop, 
+    seed,
+    sampling_steps,
+    cfg_scale):
+    return run_inference(
+        "Text2Image", 
+        prompt_text, 
+        resolution,
+        aspect_ratio,
+        length,
+        motion_strength,
+        aesthetic_score,
+        use_motion_strength,
+        use_aesthetic_score,
+        camera_motion,
+        reference_image,
+        refine_prompt,
+        fps,
+        num_loop, 
+        seed,
+        sampling_steps,
+        cfg_scale)
+
+@spaces.GPU(duration=200)
+def run_video_inference(
+    prompt_text,
+    resolution,
+    aspect_ratio,
+    length,
+    motion_strength,
+    aesthetic_score,
+    use_motion_strength,
+    use_aesthetic_score, 
+    camera_motion,
+    reference_image, 
+    refine_prompt,
+    fps,
+    num_loop, 
+    seed,
+    sampling_steps,
+    cfg_scale):
+    # if (resolution == "480p" and length == "16s") or \
+    #     (resolution == "720p" and length in ["8s", "16s"]):
+    #     gr.Warning("Generation is interrupted as the combination of 480p and 16s will lead to CUDA out of memory")
+    # else:
+    return run_inference(
+        "Text2Video",
+        prompt_text, 
+        resolution,
+        aspect_ratio, 
+        length, 
+        motion_strength, 
+        aesthetic_score, 
+        use_motion_strength,
+        use_aesthetic_score, 
+        camera_motion,
+        reference_image, 
+        refine_prompt,
+        fps,
+        num_loop, 
+        seed,
+        sampling_steps, 
+        cfg_scale
+        )
+
+
+def generate_random_prompt():
+    if "OPENAI_API_KEY" not in os.environ:
+        gr.Warning("Your prompt is empty and the OpenAI API key is not provided, please enter a valid prompt")
+        return None
+    else:
+        prompt_text = get_random_prompt_by_openai()
+        return prompt_text
 
 
 def main():
@@ -442,31 +497,119 @@ def main():
 
         with gr.Row():
             with gr.Column():
-                mode = gr.Radio(
-                    choices=["Text2Video", "Image2Video"], 
-                    value="Text2Video",
-                    label="Usage", 
-                    info="Choose your usage scenario",
-                )
                 prompt_text = gr.Textbox(
                     label="Prompt",
                     placeholder="Describe your video here",
-                    lines=4,
+                    lines=4
                 )
+                refine_prompt = gr.Checkbox(value=True, label="Refine prompt with GPT4o")
+                random_prompt_btn = gr.Button("Random Prompt By GPT4o")
+                
+                gr.Markdown("## Basic Settings")
                 resolution = gr.Radio(
-                     choices=["144p", "240p", "360p", "480p", "720p", "1080p"],
-                     value="144p",
+                     choices=["144p", "240p", "360p", "480p", "720p"],
+                     value="480p",
                     label="Resolution", 
                 )
+                aspect_ratio = gr.Radio(
+                     choices=["9:16", "16:9", "3:4", "4:3", "1:1"],
+                     value="9:16",
+                    label="Aspect Ratio (H:W)", 
+                )
                 length = gr.Radio(
-                    choices=["2s", "4s", "8s"], 
+                    choices=["2s", "4s", "8s", "16s"], 
                     value="2s",
                     label="Video Length", 
-                    info="8s may fail as Hugging Face ZeroGPU has the limitation of max 200 seconds inference time."
+                    info="only effective for video generation, 8s may fail as Hugging Face ZeroGPU has the limitation of max 200 seconds inference time."
                 )
 
+                with gr.Row():
+                    seed = gr.Slider(
+                        value=1024,
+                        minimum=1,
+                        maximum=2048,
+                        step=1,
+                        label="Seed"
+                    )
+
+                    sampling_steps = gr.Slider(
+                        value=30,
+                        minimum=1,
+                        maximum=200,
+                        step=1,
+                        label="Sampling steps"
+                    )
+                    cfg_scale = gr.Slider(
+                        value=7.0,
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.1,
+                        label="CFG Scale"
+                    )
+                
+                with gr.Row():
+                    with gr.Column():
+                        motion_strength = gr.Slider(
+                            value=5,
+                            minimum=0,
+                            maximum=100,
+                            step=1,
+                            label="Motion Strength",
+                            info="only effective for video generation"
+                        )
+                        use_motion_strength = gr.Checkbox(value=False, label="Enable")
+                        
+                    with gr.Column():
+                        aesthetic_score = gr.Slider(
+                            value=6.5,
+                            minimum=4,
+                            maximum=7,
+                            step=0.1,
+                            label="Aesthetic",
+                            info="effective for text & video generation"
+                        )
+                        use_aesthetic_score = gr.Checkbox(value=True, label="Enable")
+                        
+                camera_motion = gr.Radio(
+                    value="none",
+                    label="Camera Motion",
+                    choices=[
+                        "none",
+                        "pan right", 
+                        "pan left",
+                        "tilt up",
+                        "tilt down",
+                        "zoom in",
+                        "zoom out", 
+                        "static"
+                        ],
+                    interactive=True
+                )
+                
+                gr.Markdown("## Advanced Settings")
+                with gr.Row():
+                    fps = gr.Slider(
+                        value=24,
+                        minimum=1,
+                        maximum=60,
+                        step=1,
+                        label="FPS",
+                        info="This is the frames per seconds for video generation, keep it to 24 if you are not sure"
+                    )
+                    num_loop = gr.Slider(
+                        value=1,
+                        minimum=1,
+                        maximum=20,
+                        step=1,
+                        label="Number of Loops",
+                        info="This will change the length of the generated video, keep it to 1 if you are not sure"
+                    )
+                    
+                
+                gr.Markdown("## Reference Image")
                 reference_image = gr.Image(
-                    label="Reference Image (only used for Image2Video)",
+                    label="Image (optional)",
+                    show_download_button=True
                 )
             
             with gr.Column():
@@ -476,14 +619,24 @@ def main():
                 )
 
         with gr.Row():
-             submit_button = gr.Button("Generate video")
+             image_gen_button = gr.Button("Generate image")
+             video_gen_button = gr.Button("Generate video")
         
 
-        submit_button.click(
-             fn=run_inference, 
-             inputs=[mode, prompt_text, resolution, length, reference_image], 
+        image_gen_button.click(
+             fn=run_image_inference, 
+             inputs=[prompt_text, resolution, aspect_ratio, length, motion_strength, aesthetic_score, use_motion_strength, use_aesthetic_score, camera_motion, reference_image, refine_prompt, fps, num_loop, seed, sampling_steps, cfg_scale], 
+             outputs=reference_image
+             )
+        video_gen_button.click(
+             fn=run_video_inference, 
+             inputs=[prompt_text, resolution, aspect_ratio, length, motion_strength, aesthetic_score, use_motion_strength, use_aesthetic_score, camera_motion, reference_image, refine_prompt, fps, num_loop, seed, sampling_steps, cfg_scale], 
              outputs=output_video
              )
+        random_prompt_btn.click(
+            fn=generate_random_prompt,
+            outputs=prompt_text
+        )
 
     # launch
     demo.launch(server_port=args.port, server_name=args.host, share=args.share)

configs/dit/train/16x256x256.py

@@ -18,7 +18,7 @@ sp_size = 1
 model = dict(
     type="DiT-XL/2",
     from_pretrained="DiT-XL-2-256x256.pt",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/dit/train/1x256x256.py

@@ -19,7 +19,7 @@ sp_size = 1
 model = dict(
     type="DiT-XL/2",
     no_temporal_pos_emb=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/latte/train/16x256x256.py

@@ -17,7 +17,7 @@ sp_size = 1
 # Define model
 model = dict(
     type="Latte-XL/2",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/inference/sample-ref.py

@@ -7,33 +7,35 @@ multi_resolution = "STDiT2"
 # Condition
 prompt_path = None
 prompt = [
-    "A car driving on the ocean.",
-    'Drone view of waves crashing against the rugged cliffs along Big Sur\'s garay point beach. The crashing blue waters create white-tipped waves, while the golden light of the setting sun illuminates the rocky shore. A small island with a lighthouse sits in the distance, and green shrubbery covers the cliff\'s edge. The steep drop from the road down to the beach is a dramatic feat, with the cliff\'s edges jutting out over the sea. This is a view that captures the raw beauty of the coast and the rugged landscape of the Pacific Coast Highway.{"reference_path": "assets/images/condition/cliff.png", "mask_strategy": "0"}',
-    "In an ornate, historical hall, a massive tidal wave peaks and begins to crash. Two surfers, seizing the moment, skillfully navigate the face of the wave.",
+    'Drone view of waves crashing against the rugged cliffs along Big Sur\'s garay point beach. {"reference_path": "assets/images/condition/cliff.png", "mask_strategy": "0"}',
+    'A breathtaking sunrise scene.{"reference_path": "assets/images/condition/sunset1.png","mask_strategy": "0"}',
+    'A car driving on the ocean.{"reference_path": "https://cdn.openai.com/tmp/s/interp/d0.mp4","mask_strategy": "0,0,-8,0,8"}',
+    'A snowy forest.{"reference_path": "https://cdn.pixabay.com/video/2021/04/25/72171-542991404_large.mp4","mask_strategy": "0,0,0,0,15,0.8"}',
+    'A breathtaking sunrise scene.{"reference_path": "assets/images/condition/sunset1.png;assets/images/condition/sunset2.png","mask_strategy": "0;0,1,0,-1,1"}',
+    '|0|a white jeep equipped with a roof rack driving on a dirt road in a coniferous forest.|2|a white jeep equipped with a roof rack driving on a dirt road in the desert.|4|a white jeep equipped with a roof rack driving on a dirt road in a mountain.|6|A white jeep equipped with a roof rack driving on a dirt road in a city.|8|a white jeep equipped with a roof rack driving on a dirt road on the surface of a river.|10|a white jeep equipped with a roof rack driving on a dirt road under the lake.|12|a white jeep equipped with a roof rack flying into the sky.|14|a white jeep equipped with a roof rack driving in the universe. Earth is the background.{"reference_path": "https://cdn.openai.com/tmp/s/interp/d0.mp4", "mask_strategy": "0,0,0,0,15"}',
 ]
 
 loop = 2
 condition_frame_length = 4
-reference_path = [
-    "https://cdn.openai.com/tmp/s/interp/d0.mp4",
-    None,
-    "assets/images/condition/wave.png",
-]
-# valid when reference_path is not None
-# (loop id, ref id, ref start, length, target start)
-mask_strategy = [
-    "0,0,0,0,8,0.3",
-    None,
-    "0",
-]
+# (
+#   loop id, [the loop index of the condition image or video]
+#   reference id, [the index of the condition image or video in the reference_path]
+#   reference start, [the start frame of the condition image or video]
+#   target start, [the location to insert]
+#   length, [the number of frames to insert]
+#   edit_ratio [the edit rate of the condition image or video]
+# )
+# See https://github.com/hpcaitech/Open-Sora/blob/main/docs/config.md#advanced-inference-config for more details
+# See https://github.com/hpcaitech/Open-Sora/blob/main/docs/commands.md#inference-with-open-sora-11 for more examples
 
 # Define model
 model = dict(
     type="STDiT2-XL/2",
-    from_pretrained=None,
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v2-stage3",
     input_sq_size=512,
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/inference/sample.py

@@ -7,10 +7,11 @@ multi_resolution = "STDiT2"
 # Define model
 model = dict(
     type="STDiT2-XL/2",
-    from_pretrained=None,
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v2-stage3",
     input_sq_size=512,
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/benchmark.py

@@ -65,7 +65,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/image.py

@@ -29,7 +29,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/image_rflow.py

@@ -0,0 +1,88 @@
+# Define dataset
+# dataset = dict(
+#     type="VariableVideoTextDataset",
+#     data_path=None,
+#     num_frames=None,
+#     frame_interval=3,
+#     image_size=(None, None),
+#     transform_name="resize_crop",
+# )
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=1,
+    frame_interval=1,
+    image_size=(256, 256),
+    transform_name="center",
+)
+bucket_config = {  # 6s/it
+    "256": {1: (1.0, 256)},
+    "512": {1: (1.0, 80)},
+    "480p": {1: (1.0, 52)},
+    "1024": {1: (1.0, 20)},
+    "1080p": {1: (1.0, 8)},
+}
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+# model = dict(
+#     type="DiT-XL/2",
+#     from_pretrained="/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/PixArt-XL-2-512x512.pth",
+#     # input_sq_size=512,  # pretrained model is trained on 512x512
+#     enable_flash_attn=True,
+#     enable_layernorm_kernel=True,
+# )
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PixArt-XL-2-512x512.pth",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+# model = dict(
+#     type="DiT-XL/2",
+#     # space_scale=1.0,
+#     # time_scale=1.0,
+#     no_temporal_pos_emb=True,
+#     # from_pretrained="PixArt-XL-2-512x512.pth",
+#     from_pretrained="/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/PixArt-XL-2-512x512.pth",
+#     enable_flash_attn=True,
+#     enable_layernorm_kernel=True,
+# )
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=200,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 10
+log_every = 10
+ckpt_every = 500
+load = None
+
+batch_size = 100  # only for logging
+lr = 2e-5
+grad_clip = 1.0

configs/opensora-v1-1/train/stage1.py

@@ -16,15 +16,15 @@ bucket_config = {  # 1s/it
     "1024": {1: (0.3, 3)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -41,7 +41,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/stage2.py

@@ -18,15 +18,15 @@ bucket_config = {  # 7s/it
     "1080p": {1: (0.4, 8)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -43,7 +43,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/stage3.py

@@ -18,15 +18,15 @@ bucket_config = {  # 13s/it
     "1024": {1: (0.3, 40)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -43,7 +43,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/video.py

@@ -31,7 +31,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-2/inference/sample.py

@@ -0,0 +1,43 @@
+resolution = "240p"
+aspect_ratio = "9:16"
+num_frames = 51
+fps = 24
+frame_interval = 1
+save_fps = 24
+
+save_dir = "./samples/samples/"
+seed = 42
+batch_size = 1
+multi_resolution = "STDiT2"
+dtype = "bf16"
+condition_frame_length = 5
+align = 5
+
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v3",
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="hpcai-tech/OpenSora-VAE-v1.2",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    num_sampling_steps=30,
+    cfg_scale=7.0,
+)
+
+aes = 6.5
+flow = None

configs/opensora-v1-2/misc/bs.py

@@ -0,0 +1,117 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# == Config 1: Webvid ==
+# base: (512, 408), 12s/it
+grad_checkpoint = True
+base = ("512", "408")
+base_step_time = 12
+bucket_config = {
+    "144p": {
+        1: (475, 0),
+        51: (51, 0),
+        102: (27, 0),
+        204: (13, 0),
+        408: (6, 0),
+    },
+    # ---
+    "240p": {
+        1: (297, 200),  # 8.25
+        51: (20, 0),
+        102: (10, 0),
+        204: (5, 0),
+        408: (2, 0),
+    },
+    # ---
+    "512": {
+        1: (141, 0),
+        51: (8, 0),
+        102: (4, 0),
+        204: (2, 0),
+        408: (1, 0),
+    },
+    # ---
+    "480p": {
+        1: (89, 0),
+        51: (5, 0),
+        102: (2, 0),
+        204: (1, 0),
+    },
+    # ---
+    "1024": {
+        1: (36, 0),
+        51: (1, 0),
+    },
+    # ---
+    "1080p": {1: (5, 0)},
+    # ---
+    "2048": {1: (5, 0)},
+}
+
+# == Config 1 ==
+# base: (512, 408), 16s/it
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/misc/eval_loss.py

@@ -0,0 +1,49 @@
+num_workers = 8
+dtype = "bf16"
+seed = 42
+num_eval_timesteps = 10
+
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+bucket_config = {
+    "144p": {1: (None, 100), 51: (None, 30), 102: (None, 20), 204: (None, 8), 408: (None, 4)},
+    # ---
+    "240p": {1: (None, 100), 51: (None, 24), 102: (None, 12), 204: (None, 4), 408: (None, 2)},
+    # ---
+    "360p": {1: (None, 60), 51: (None, 12), 102: (None, 6), 204: (None, 2), 408: (None, 1)},
+    # ---
+    "480p": {1: (None, 40), 51: (None, 6), 102: (None, 3), 204: (None, 1)},
+    # ---
+    "720p": {1: (None, 20), 51: (None, 2), 102: (None, 1)},
+    # ---
+    "1080p": {1: (None, 10)},
+    # ---
+    "2048": {1: (None, 5)},
+}
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+    local_files_only=True,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    local_files_only=True,
+)
+scheduler = dict(type="rflow")

configs/opensora-v1-2/misc/extract.py

@@ -0,0 +1,62 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+seed = 42
+outputs = "outputs"
+wandb = False
+
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained="/mnt/nfs-206/zangwei/opensora/outputs/1091-STDiT3-XL-2/epoch0-global_step8500",
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=32,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+
+# feature extraction settings
+save_text_features = True
+save_compressed_text_features = True
+bin_size = 250  # 1GB, 4195 bins
+log_time = False

configs/opensora-v1-2/misc/feat.py

@@ -0,0 +1,94 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+    dummy_text_feature=True,
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+grad_checkpoint = True
+
+load_text_features = True
+
+# Acceleration settings
+num_workers = 0
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+    skip_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 1
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/adapt.py

@@ -0,0 +1,84 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+bucket_config = {  # 2s/it
+    "144p": {1: (0.5, 48), 34: (1.0, 2), 51: (1.0, 4), 102: (1.0, 2), 204: (1.0, 1)},
+    # ---
+    "256": {1: (0.6, 20), 34: (0.5, 2), 51: (0.5, 1), 68: (0.5, 1), 136: (0.0, None)},
+    "240p": {1: (0.6, 20), 34: (0.5, 2), 51: (0.5, 1), 68: (0.5, 1), 136: (0.0, None)},
+    # ---
+    "360p": {1: (0.5, 8), 34: (0.2, 1), 102: (0.0, None)},
+    "512": {1: (0.5, 8), 34: (0.2, 1), 102: (0.0, None)},
+    # ---
+    "480p": {1: (0.2, 4), 17: (0.3, 1), 68: (0.0, None)},
+    # ---
+    "720p": {1: (0.1, 2)},
+    "1024": {1: (0.1, 2)},
+    # ---
+    "1080p": {1: (0.1, 1)},
+}
+grad_checkpoint = False
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/stage1.py

@@ -0,0 +1,111 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# backup
+# bucket_config = {  # 20s/it
+#     "144p": {1: (1.0, 100), 51: (1.0, 30), 102: (1.0, 20), 204: (1.0, 8), 408: (1.0, 4)},
+#     # ---
+#     "256": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
+#     "240p": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
+#     # ---
+#     "360p": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
+#     "512": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
+#     # ---
+#     "480p": {1: (0.5, 40), 51: (0.3, 6), 102: (0.3, 3), 204: (0.3, 1), 408: (0.0, None)},
+#     # ---
+#     "720p": {1: (0.2, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
+#     "1024": {1: (0.1, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
+#     # ---
+#     "1080p": {1: (0.1, 10)},
+#     # ---
+#     "2048": {1: (0.1, 5)},
+# }
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.05,
+    "intepolate": 0.005,
+    "quarter_random": 0.005,
+    "quarter_head": 0.005,
+    "quarter_tail": 0.005,
+    "quarter_head_tail": 0.005,
+    "image_random": 0.025,
+    "image_head": 0.05,
+    "image_tail": 0.025,
+    "image_head_tail": 0.025,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15
+warmup_steps = 1000

configs/opensora-v1-2/train/stage1_feat.py

@@ -0,0 +1,59 @@
+# Dataset settings
+dataset = dict(type="BatchFeatureDataset")
+grad_checkpoint = True
+num_workers = 4
+
+# Acceleration settings
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+    skip_y_embedder=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+vae_out_channels = 4
+model_max_length = 300
+text_encoder_output_dim = 4096
+load_video_features = True
+load_text_features = True
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/stage2.py

@@ -0,0 +1,91 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 1.0), 5), 408: ((0.5, 1.0), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 1.0), 5), 408: ((0.4, 1.0), 2)},
+    # ---
+    "360p": {1: (0.5, 141), 51: (0.15, 8), 102: ((0.3, 0.5), 4), 204: ((0.3, 1.0), 2), 408: ((0.5, 0.5), 1)},
+    "512": {1: (0.4, 141), 51: (0.15, 8), 102: ((0.2, 0.4), 4), 204: ((0.2, 1.0), 2), 408: ((0.4, 0.5), 1)},
+    # ---
+    "480p": {1: (0.5, 89), 51: (0.2, 5), 102: (0.2, 2), 204: (0.1, 1)},
+    # ---
+    "720p": {1: (0.1, 36), 51: (0.03, 1)},
+    "1024": {1: (0.1, 36), 51: (0.02, 1)},
+    # ---
+    "1080p": {1: (0.01, 5)},
+    # ---
+    "2048": {1: (0.01, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.05,
+    "intepolate": 0.005,
+    "quarter_random": 0.005,
+    "quarter_head": 0.005,
+    "quarter_tail": 0.005,
+    "quarter_head_tail": 0.005,
+    "image_random": 0.025,
+    "image_head": 0.05,
+    "image_tail": 0.025,
+    "image_head_tail": 0.025,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15
+

configs/opensora-v1-2/train/stage3.py

@@ -0,0 +1,92 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 20s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: (1.0, 27), 204: (1.0, 13), 408: (1.0, 6)},
+    # ---
+    "256": {1: (1.0, 297), 51: (0.5, 20), 102: (0.5, 10), 204: (0.5, 5), 408: ((0.5, 0.5), 2)},
+    "240p": {1: (1.0, 297), 51: (0.5, 20), 102: (0.5, 10), 204: (0.5, 5), 408: ((0.5, 0.4), 2)},
+    # ---
+    "360p": {1: (1.0, 141), 51: (0.5, 8), 102: (0.5, 4), 204: (0.5, 2), 408: ((0.5, 0.3), 1)},
+    "512": {1: (1.0, 141), 51: (0.5, 8), 102: (0.5, 4), 204: (0.5, 2), 408: ((0.5, 0.2), 1)},
+    # ---
+    "480p": {1: (1.0, 89), 51: (0.5, 5), 102: (0.5, 3), 204: ((0.5, 0.5), 1), 408: (0.0, None)},
+    # ---
+    "720p": {1: (0.3, 36), 51: (0.2, 2), 102: (0.1, 1), 204: (0.0, None)},
+    "1024": {1: (0.3, 36), 51: (0.1, 2), 102: (0.1, 1), 204: (0.0, None)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.05, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+# 25%
+mask_ratios = {
+    "random": 0.01,
+    "intepolate": 0.002,
+    "quarter_random": 0.002,
+    "quarter_head": 0.002,
+    "quarter_tail": 0.002,
+    "quarter_head_tail": 0.002,
+    "image_random": 0.0,
+    "image_head": 0.22,
+    "image_tail": 0.005,
+    "image_head_tail": 0.005,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15
+warmup_steps = 1000

configs/opensora/inference/16x256x256.py

@@ -7,7 +7,7 @@ model = dict(
     type="STDiT-XL/2",
     space_scale=0.5,
     time_scale=1.0,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

configs/opensora/inference/16x512x512-rflow.py

@@ -0,0 +1,35 @@
+num_frames = 16
+fps = 24 // 3
+image_size = (512, 512)
+
+# Define model
+model = dict(
+    type="STDiT-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    from_pretrained="PRETRAINED_MODEL",
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+    micro_batch_size=2,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+)
+scheduler = dict(
+    type="rflow",
+    num_sampling_steps=10,
+    cfg_scale=7.0,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 2
+seed = 42
+prompt_path = "./assets/texts/t2v_samples.txt"
+save_dir = "./outputs/samples/"

configs/opensora/inference/16x512x512.py

@@ -7,7 +7,7 @@ model = dict(
     type="STDiT-XL/2",
     space_scale=1.0,
     time_scale=1.0,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

configs/opensora/inference/64x512x512.py

@@ -7,7 +7,7 @@ model = dict(
     type="STDiT-XL/2",
     space_scale=1.0,
     time_scale=2 / 3,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

configs/opensora/train/16x256x256-mask.py

@@ -20,12 +20,12 @@ model = dict(
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 mask_ratios = {
-    "mask_no": 0.7,
-    "mask_random": 0.15,
+    "identity": 0.7,
+    "random": 0.15,
     "mask_head": 0.05,
     "mask_tail": 0.05,
     "mask_head_tail": 0.05,

configs/opensora/train/16x256x256-spee-rflow.py

@@ -0,0 +1,64 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=16,
+    frame_interval=3,
+    image_size=(256, 256),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="STDiT-XL/2",
+    space_scale=0.5,
+    time_scale=1.0,
+    # from_pretrained="PixArt-XL-2-512x512.pth",
+    # from_pretrained = "/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/OpenSora-v1-HQ-16x512x512.pth",
+    # from_pretrained = "OpenSora-v1-HQ-16x512x512.pth",
+    from_pretrained="PRETRAINED_MODEL",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+# mask_ratios = [0.5, 0.29, 0.07, 0.07, 0.07]
+# mask_ratios = {
+#     "identity": 0.9,
+#     "random": 0.06,
+#     "mask_head": 0.01,
+#     "mask_tail": 0.01,
+#     "mask_head_tail": 0.02,
+# }
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = True
+
+epochs = 1
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 16
+lr = 2e-5
+grad_clip = 1.0

configs/opensora/train/16x256x256-spee.py

@@ -20,12 +20,12 @@ model = dict(
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 mask_ratios = {
-    "mask_no": 0.5,
-    "mask_random": 0.29,
+    "identity": 0.5,
+    "random": 0.29,
     "mask_head": 0.07,
     "mask_tail": 0.07,
     "mask_head_tail": 0.07,

configs/opensora/train/16x256x256.py

@@ -8,7 +8,7 @@ dataset = dict(
 )
 
 # Define acceleration
-num_workers = 4
+num_workers = 0
 dtype = "bf16"
 grad_checkpoint = True
 plugin = "zero2"
@@ -20,7 +20,7 @@ model = dict(
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora/train/16x512x512.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=1.0,
     time_scale=1.0,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora/train/360x512x512.py

@@ -26,7 +26,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     enable_sequence_parallelism=True,  # enable sq here
 )

configs/opensora/train/64x512x512-sp.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     enable_sequence_parallelism=True,  # enable sq here
 )

configs/opensora/train/64x512x512.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/pixart/inference/1x20481B.py

@@ -0,0 +1,36 @@
+num_frames = 1
+fps = 1
+image_size = (2560, 1536)
+# image_size = (2048, 2048)
+
+model = dict(
+    type="PixArt-1B/2",
+    from_pretrained="PixArt-1B-2.pth",
+    space_scale=4,
+    no_temporal_pos_emb=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    base_size=2048 // 8,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="dpm-solver",
+    num_sampling_steps=14,
+    cfg_scale=4.5,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 1
+seed = 42
+prompt_path = "./assets/texts/t2i_sigma.txt"
+save_dir = "./samples/samples/"

configs/pixart/inference/1x2048MS.py

@@ -0,0 +1,36 @@
+num_frames = 1
+fps = 1
+image_size = (2560, 1536)
+# image_size = (2048, 2048)
+
+model = dict(
+    type="PixArt-XL/2",
+    from_pretrained="PixArt-Sigma-XL-2-2K-MS.pth",
+    space_scale=4,
+    no_temporal_pos_emb=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    base_size=2048 // 8,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="dpm-solver",
+    num_sampling_steps=14,
+    cfg_scale=4.5,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 1
+seed = 42
+prompt_path = "./assets/texts/t2i_sigma.txt"
+save_dir = "./samples/samples/"

configs/pixart/inference/1x512x512-rflow.py

@@ -0,0 +1,39 @@
+num_frames = 1
+fps = 1
+image_size = (512, 512)
+
+# Define model
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PRETRAINED_MODEL",
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+)
+scheduler = dict(
+    type="rflow",
+    num_sampling_steps=20,
+    cfg_scale=7.0,
+)
+dtype = "bf16"
+
+# prompt_path = "./assets/texts/t2i_samples.txt"
+prompt = [
+    "Pirate ship trapped in a cosmic maelstrom nebula.",
+    "A small cactus with a happy face in the Sahara desert.",
+    "A small cactus with a sad face in the Sahara desert.",
+]
+
+# Others
+batch_size = 2
+seed = 42
+save_dir = "./outputs/samples2/"

configs/pixart/train/16x256x256.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/pixart/train/1x2048x2048.py

@@ -0,0 +1,54 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path="/home/zhaowangbo/data/csv/image-v1_1_ext_noempty_rcp_clean_info.csv",
+    num_frames=1,
+    frame_interval=3,
+    image_size=(2048, 2048),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="PixArt-1B/2",
+    space_scale=4.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PixArt-1B-2.pth",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="iddpm",
+    timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 1000
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 4
+lr = 2e-5
+grad_clip = 1.0

configs/pixart/train/1x512x512-rflow.py

@@ -0,0 +1,55 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=1,
+    frame_interval=3,
+    image_size=(512, 512),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    # from_pretrained="PixArt-XL-2-512x512.pth",
+    from_pretrained="PRETRAINED_MODEL",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = True
+
+epochs = 2
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 64
+lr = 2e-5
+grad_clip = 1.0

configs/pixart/train/1x512x512.py

@@ -21,7 +21,7 @@ model = dict(
     time_scale=1.0,
     no_temporal_pos_emb=True,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/pixart/train/64x512x512.py

@@ -21,7 +21,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/vae/inference/image.py

@@ -0,0 +1,32 @@
+image_size = (256, 256)
+num_frames = 1
+
+dtype = "bf16"
+batch_size = 1
+seed = 42
+save_dir = "samples/vae_video"
+cal_stats = True
+log_stats_every = 100
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    image_size=image_size,
+)
+num_samples = 100
+num_workers = 4
+
+# Define model
+model = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=None,
+    micro_batch_size=4,
+    cal_loss=True,
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6

configs/vae/inference/video.py

@@ -0,0 +1,32 @@
+image_size = (256, 256)
+num_frames = 17
+
+dtype = "bf16"
+batch_size = 1
+seed = 42
+save_dir = "samples/vae_video"
+cal_stats = True
+log_stats_every = 100
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    image_size=image_size,
+)
+num_samples = 100
+num_workers = 4
+
+# Define model
+model = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=None,
+    micro_batch_size=4,
+    cal_loss=True,
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6

configs/vae/train/image.py

@@ -0,0 +1,58 @@
+num_frames = 1
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=True,
+    from_pretrained=None,
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+# loss weights
+perceptual_loss_weight = 0.0  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_image_ratio = 0.1
+use_real_rec_loss = False
+use_z_rec_loss = True
+use_image_identity_loss = True
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 100
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0

configs/vae/train/video.py

@@ -0,0 +1,58 @@
+num_frames = 33
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=False,
+    from_pretrained=None,
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_image_ratio = 0.2
+use_real_rec_loss = True
+use_z_rec_loss = False
+use_image_identity_loss = False
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 100
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0

configs/vae/train/video_disc.py

@@ -0,0 +1,75 @@
+num_frames = 17
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=False,
+    from_pretrained=None,
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+discriminator = dict(
+    type="NLayerDiscriminator",
+    from_pretrained="/home/shenchenhui/opensoraplan-v1.0.0-discriminator.pt",
+    input_nc=3,
+    n_layers=3,
+    use_actnorm=False,
+)
+
+# discriminator hyper-parames TODO
+discriminator_factor = 1
+discriminator_start = -1
+generator_factor = 0.5
+generator_loss_type = "hinge"
+discriminator_loss_type = "hinge"
+lecam_loss_weight = None
+gradient_penalty_loss_weight = None
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_image_ratio = 0.2
+use_real_rec_loss = True
+use_z_rec_loss = False
+use_image_identity_loss = False
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 100
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0

requirements.txt

@@ -1,4 +1,4 @@
 xformers
 transformers
 pandarallel
-git+https://github.com/hpcaitech/Open-Sora.git#egg=opensora
+git+https://github.com/FrankLeeeee/test-space.git