| # ========================================================= |
| # ZERO GPU PATCHED + ALL TASKS ENABLED |
| # Hugging Face Spaces Compatible |
| # ========================================================= |
|
|
| from __future__ import annotations |
|
|
| import gc |
| import json |
| import os |
| import random |
| import threading |
| import time |
| import traceback |
|
|
| from copy import deepcopy |
| from pathlib import Path |
|
|
| import gradio as gr |
| import spaces |
| import torch |
|
|
| # ========================================================= |
| # ENV |
| # ========================================================= |
|
|
| os.environ["TOKENIZERS_PARALLELISM"] = "false" |
| os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" |
| os.environ["TRANSFORMERS_NO_FLASH_ATTN"] = "1" |
| os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" |
| os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" |
|
|
| # ========================================================= |
| # LOGIN |
| # ========================================================= |
|
|
| from huggingface_hub import login, snapshot_download |
|
|
| HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
| if HF_TOKEN: |
| login(token=HF_TOKEN) |
|
|
| # ========================================================= |
| # IMPORTS |
| # ========================================================= |
|
|
| from safetensors.torch import load_file |
| from transformers import set_seed |
| from transformers.utils import is_flash_attn_2_available |
|
|
| from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( |
| Qwen2_5_VLVisionConfig, |
| ) |
|
|
| from common.utils.logging import get_logger |
|
|
| from config.config_factory import ( |
| DataArguments, |
| InferenceArguments, |
| ModelArguments, |
| ) |
|
|
| from data.data_utils import add_special_tokens |
| from data.dataset_base import DataConfig, simple_custom_collate |
| from data.datasets_custom import ValidationDataset |
|
|
| from inference_lance import ( |
| apply_inference_defaults, |
| clean_memory, |
| init_from_model_path_if_needed, |
| validate_on_fixed_batch, |
| ) |
|
|
| from modeling.lance import ( |
| Lance, |
| LanceConfig, |
| Qwen2ForCausalLM, |
| ) |
|
|
| from modeling.qwen2 import Qwen2Tokenizer |
| from modeling.qwen2.modeling_qwen2 import Qwen2Config |
|
|
| from modeling.vae.wan.model import WanVideoVAE |
|
|
| from modeling.vit.qwen2_5_vl_vit import ( |
| Qwen2_5_VisionTransformerPretrainedModel, |
| ) |
|
|
| print("FlashAttention2 available:", is_flash_attn_2_available()) |
|
|
| # ========================================================= |
| # PERFORMANCE |
| # ========================================================= |
|
|
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
| # ========================================================= |
| # PATHS |
| # ========================================================= |
|
|
| REPO_ROOT = Path(__file__).resolve().parent |
|
|
| TMP_ROOT = REPO_ROOT / "tmps" |
| TMP_ROOT.mkdir(parents=True, exist_ok=True) |
|
|
| TMP_INPUT_DIR = TMP_ROOT / "inputs" |
| TMP_INPUT_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| RESULTS_ROOT = TMP_ROOT / "results" |
| RESULTS_ROOT.mkdir(parents=True, exist_ok=True) |
|
|
| # ========================================================= |
| # MODEL DOWNLOAD |
| # ========================================================= |
|
|
| MODEL_REPO = "bytedance-research/Lance" |
|
|
| MODEL_CACHE_DIR = REPO_ROOT / "downloads" |
|
|
| snapshot_download( |
| repo_id=MODEL_REPO, |
| allow_patterns=[ |
| "Lance_3B_Video/*", |
| ], |
| local_dir=str(MODEL_CACHE_DIR), |
| local_dir_use_symlinks=False, |
| token=HF_TOKEN, |
| ) |
|
|
| DEFAULT_MODEL_PATH = str( |
| MODEL_CACHE_DIR / "Lance_3B_Video" |
| ) |
|
|
| print("DEFAULT_MODEL_PATH =", DEFAULT_MODEL_PATH) |
|
|
| # ========================================================= |
| # DEFAULTS |
| # ========================================================= |
|
|
| DEFAULT_TIMESTEPS = 30 |
| DEFAULT_TIMESTEP_SHIFT = 3.5 |
| DEFAULT_CFG_TEXT_SCALE = 4.0 |
|
|
| DEFAULT_HEIGHT = 480 |
| DEFAULT_WIDTH = 848 |
| DEFAULT_NUM_FRAMES = 50 |
|
|
| DEFAULT_RESOLUTION = "video_480p" |
|
|
| DEFAULT_TASK = "t2v" |
|
|
| DEFAULT_SEED = -1 |
|
|
| DEFAULT_VIT_TYPE = "qwen_2_5_vl_original" |
|
|
| # ========================================================= |
| # TASKS |
| # ========================================================= |
|
|
| TASK_T2V = "t2v" |
| TASK_T2I = "t2i" |
|
|
| TASK_IMAGE_EDIT = "image_edit" |
| TASK_VIDEO_EDIT = "video_edit" |
|
|
| TASK_X2T_IMAGE = "x2t_image" |
| TASK_X2T_VIDEO = "x2t_video" |
|
|
| TASK_CHOICES = [ |
| TASK_T2V, |
| TASK_T2I, |
| TASK_IMAGE_EDIT, |
| TASK_VIDEO_EDIT, |
| TASK_X2T_IMAGE, |
| TASK_X2T_VIDEO, |
| ] |
|
|
| VIDEO_RESOLUTION_CHOICES = [ |
| "video_192p", |
| "video_360p", |
| "video_480p", |
| ] |
|
|
| # ========================================================= |
| # HELPERS |
| # ========================================================= |
|
|
| def normalize_seed(seed: int): |
|
|
| if seed == -1: |
| return random.randint(0, 2**31 - 1) |
|
|
| return seed |
|
|
|
|
| def normalize_task(task: str): |
|
|
| task = (task or DEFAULT_TASK).strip().lower() |
|
|
| valid = { |
| TASK_T2V, |
| TASK_T2I, |
| TASK_IMAGE_EDIT, |
| TASK_VIDEO_EDIT, |
| TASK_X2T_IMAGE, |
| TASK_X2T_VIDEO, |
| } |
|
|
| if task not in valid: |
| return DEFAULT_TASK |
|
|
| return task |
|
|
| # ========================================================= |
| # PIPELINE |
| # ========================================================= |
| |
| class LancePipeline: |
| |
| def __init__(self): |
| |
| self.initialized = False |
| |
| self.lock = threading.Lock() |
| |
| self.logger = get_logger("lance") |
| |
| def initialize(self): |
| |
| with self.lock: |
| |
| if self.initialized: |
| return |
| |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA unavailable") |
| |
| model_args = ModelArguments( |
| model_path=DEFAULT_MODEL_PATH, |
| vit_type=DEFAULT_VIT_TYPE, |
| llm_qk_norm=True, |
| llm_qk_norm_und=True, |
| llm_qk_norm_gen=True, |
| tie_word_embeddings=False, |
| max_num_frames=121, |
| max_latent_size=64, |
| latent_patch_size=[1,1,1], |
| ) |
| |
| data_args = DataArguments() |
| |
| inference_args = InferenceArguments( |
| validation_num_timesteps=DEFAULT_TIMESTEPS, |
| validation_timestep_shift=DEFAULT_TIMESTEP_SHIFT, |
| copy_init_moe=True, |
| visual_und=True, |
| visual_gen=True, |
| vae_model_type="wan", |
| apply_qwen_2_5_vl_pos_emb=True, |
| apply_chat_template=False, |
| cfg_type=0, |
| validation_data_seed=42, |
| video_height=DEFAULT_HEIGHT, |
| video_width=DEFAULT_WIDTH, |
| num_frames=DEFAULT_NUM_FRAMES, |
| task=DEFAULT_TASK, |
| save_path_gen=str(RESULTS_ROOT), |
| resolution=DEFAULT_RESOLUTION, |
| text_template=True, |
| use_KVcache=True, |
| ) |
| |
| apply_inference_defaults( |
| model_args, |
| data_args, |
| inference_args, |
| ) |
| |
| set_seed(42) |
| |
| llm_config = Qwen2Config.from_json_file( |
| str(Path(model_args.model_path) / "llm_config.json") |
| ) |
| |
| language_model = Qwen2ForCausalLM(llm_config) |
| |
| vit_config = Qwen2_5_VLVisionConfig.from_pretrained( |
| model_args.vit_path |
| ) |
| |
| vit_config._attn_implementation = "eager" |
| |
| vit_model = Qwen2_5_VisionTransformerPretrainedModel( |
| vit_config |
| ) |
| |
| vit_weights = load_file( |
| str(Path(model_args.vit_path) / "vit.safetensors") |
| ) |
| |
| vit_model.load_state_dict(vit_weights, strict=True) |
| |
| clean_memory(vit_weights) |
| |
| vae_model = WanVideoVAE() |
| |
| vae_config = deepcopy(vae_model.vae_config) |
| |
| config = LanceConfig( |
| visual_gen=True, |
| visual_und=True, |
| llm_config=llm_config, |
| vit_config=vit_config, |
| vae_config=vae_config, |
| latent_patch_size=model_args.latent_patch_size, |
| max_num_frames=model_args.max_num_frames, |
| max_latent_size=model_args.max_latent_size, |
| vit_max_num_patch_per_side=model_args.vit_max_num_patch_per_side, |
| connector_act=model_args.connector_act, |
| interpolate_pos=model_args.interpolate_pos, |
| timestep_shift=inference_args.timestep_shift, |
| ) |
| |
| model = Lance( |
| language_model=language_model, |
| vit_model=vit_model, |
| vit_type=model_args.vit_type, |
| config=config, |
| training_args=inference_args, |
| ) |
| |
| model = model.to( |
| device="cuda", |
| dtype=torch.bfloat16, |
| ) |
| |
| tokenizer = Qwen2Tokenizer.from_pretrained( |
| model_args.model_path |
| ) |
| |
| tokenizer, new_token_ids, _ = add_special_tokens( |
| tokenizer |
| ) |
| |
| init_from_model_path_if_needed( |
| model, |
| model_args, |
| ) |
| |
| image_token_id = model.language_model.config.video_token_id |
| |
| model.eval() |
| |
| self.model = model |
| self.vae_model = vae_model |
| self.vae_config = vae_config |
| |
| self.tokenizer = tokenizer |
| self.new_token_ids = new_token_ids |
| self.image_token_id = image_token_id |
| |
| self.base_model_args = model_args |
| self.base_data_args = data_args |
| self.base_inference_args = inference_args |
| |
| self.initialized = True |
| |
| print("Lance initialized successfully") |
| |
| def generate( |
| self, |
| task, |
| prompt, |
| input_image, |
| input_video, |
| question, |
| height, |
| width, |
| num_frames, |
| seed, |
| resolution, |
| validation_num_timesteps, |
| validation_timestep_shift, |
| cfg_text_scale, |
| ): |
| |
| task = normalize_task(task) |
| |
| actual_seed = normalize_seed(int(seed)) |
| |
| set_seed(actual_seed) |
| |
| save_dir = RESULTS_ROOT / str(time.time()) |
| save_dir.mkdir(parents=True, exist_ok=True) |
| |
| inference_args = deepcopy( |
| self.base_inference_args |
| ) |
| |
| inference_args.video_height = int(height) |
| inference_args.video_width = int(width) |
| inference_args.num_frames = int(num_frames) |
| |
| inference_args.validation_num_timesteps = ( |
| validation_num_timesteps |
| ) |
| |
| inference_args.validation_timestep_shift = ( |
| validation_timestep_shift |
| ) |
| |
| inference_args.task = task |
| |
| prompt_file = TMP_INPUT_DIR / "prompt.json" |
| |
| # ===================================================== |
| # PAYLOADS |
| # ===================================================== |
| |
| if task == TASK_T2V: |
| |
| payload = { |
| "000000.mp4": prompt |
| } |
| |
| elif task == TASK_T2I: |
| |
| payload = { |
| "000000.png": prompt |
| } |
| |
| elif task == TASK_IMAGE_EDIT: |
| |
| payload = { |
| "000000": { |
| "interleave_array": [ |
| input_image, |
| [prompt, ""] |
| ], |
| "element_dtype_array": [ |
| "image", |
| "text" |
| ], |
| "istarget_in_interleave": [ |
| 0, |
| 1 |
| ], |
| } |
| } |
| |
| elif task == TASK_VIDEO_EDIT: |
| |
| payload = { |
| "000000": { |
| "interleave_array": [ |
| input_video, |
| [prompt, ""] |
| ], |
| "element_dtype_array": [ |
| "video", |
| "text" |
| ], |
| "istarget_in_interleave": [ |
| 0, |
| 1 |
| ], |
| } |
| } |
| |
| elif task == TASK_X2T_IMAGE: |
| |
| payload = { |
| "000000": { |
| "interleave_array": [ |
| input_image, |
| [ |
| "Describe the image", |
| question, |
| "" |
| ] |
| ], |
| "element_dtype_array": [ |
| "image", |
| "text" |
| ], |
| "istarget_in_interleave": [ |
| 0, |
| 1 |
| ], |
| } |
| } |
| |
| elif task == TASK_X2T_VIDEO: |
| |
| payload = { |
| "000000": { |
| "interleave_array": [ |
| input_video, |
| [ |
| "Describe the video", |
| question, |
| "" |
| ] |
| ], |
| "element_dtype_array": [ |
| "video", |
| "text" |
| ], |
| "istarget_in_interleave": [ |
| 0, |
| 1 |
| ], |
| } |
| } |
| |
| else: |
| |
| return ( |
| None, |
| None, |
| "", |
| "Invalid task", |
| "", |
| ) |
| |
| with open(prompt_file, "w") as f: |
| json.dump(payload, f) |
| |
| dataset_config = DataConfig.from_yaml( |
| str(prompt_file) |
| ) |
| |
| val_dataset = ValidationDataset( |
| jsonl_path=str(prompt_file), |
| tokenizer=self.tokenizer, |
| data_args=self.base_data_args, |
| model_args=self.base_model_args, |
| training_args=inference_args, |
| new_token_ids=self.new_token_ids, |
| dataset_config=dataset_config, |
| local_rank=0, |
| world_size=1, |
| ) |
| |
| val_data_cpu = simple_custom_collate( |
| [val_dataset[0]] |
| ) |
| |
| validate_on_fixed_batch( |
| fsdp_model=self.model, |
| vae_model=self.vae_model, |
| tokenizer=self.tokenizer, |
| val_data_cpu=val_data_cpu, |
| training_args=inference_args, |
| model_args=self.base_model_args, |
| inference_args=inference_args, |
| new_token_ids=self.new_token_ids, |
| image_token_id=self.image_token_id, |
| device="cuda", |
| save_source_video=False, |
| save_path_gen=str(save_dir), |
| save_path_gt="", |
| ) |
| |
| clean_memory() |
| |
| gc.collect() |
| |
| torch.cuda.empty_cache() |
| |
| videos = list(save_dir.glob("*.mp4")) |
| images = list(save_dir.glob("*.png")) |
| |
| if len(videos) > 0: |
| |
| return ( |
| str(videos[0]), |
| None, |
| "", |
| "Success", |
| "", |
| ) |
| |
| if len(images) > 0: |
| |
| return ( |
| None, |
| str(images[0]), |
| "", |
| "Success", |
| "", |
| ) |
| |
| if task in [TASK_X2T_IMAGE, TASK_X2T_VIDEO]: |
| |
| return ( |
| None, |
| None, |
| "Understanding complete", |
| "Success", |
| "", |
| ) |
| |
| return ( |
| None, |
| None, |
| "", |
| "No output generated", |
| "", |
| ) |
| |
| # ========================================================= |
| # GLOBAL |
| # ========================================================= |
| |
| PIPELINE = LancePipeline() |
| |
| # ========================================================= |
| # GPU RUN |
| # ========================================================= |
| |
| @spaces.GPU(duration=300) |
| def run_task( |
| task, |
| prompt, |
| input_image, |
| input_video, |
| question, |
| height, |
| width, |
| num_frames, |
| seed, |
| resolution, |
| validation_num_timesteps, |
| validation_timestep_shift, |
| cfg_text_scale, |
| ): |
| |
| PIPELINE.initialize() |
| |
| return PIPELINE.generate( |
| task=task, |
| prompt=prompt, |
| input_image=input_image, |
| input_video=input_video, |
| question=question, |
| height=height, |
| width=width, |
| num_frames=num_frames, |
| seed=seed, |
| resolution=resolution, |
| validation_num_timesteps=validation_num_timesteps, |
| validation_timestep_shift=validation_timestep_shift, |
| cfg_text_scale=cfg_text_scale, |
| ) |
| |
| # ========================================================= |
| # UI |
| # ========================================================= |
| |
| with gr.Blocks(title="Lance ZeroGPU") as demo: |
| |
| gr.Markdown("# Lance Multi-Task ZeroGPU") |
| |
| with gr.Row(): |
| |
| with gr.Column(): |
| |
| task = gr.Dropdown( |
| label="Task", |
| choices=TASK_CHOICES, |
| value=DEFAULT_TASK, |
| ) |
| |
| prompt = gr.Textbox( |
| label="Prompt", |
| lines=5, |
| ) |
| |
| input_image = gr.Image( |
| label="Input Image", |
| type="filepath", |
| ) |
| |
| input_video = gr.Video( |
| label="Input Video", |
| ) |
| |
| question = gr.Textbox( |
| label="Question", |
| lines=3, |
| ) |
| |
| height = gr.Slider( |
| minimum=192, |
| maximum=1024, |
| step=16, |
| value=DEFAULT_HEIGHT, |
| label="Height", |
| ) |
| |
| width = gr.Slider( |
| minimum=192, |
| maximum=1024, |
| step=16, |
| value=DEFAULT_WIDTH, |
| label="Width", |
| ) |
| |
| num_frames = gr.Slider( |
| minimum=1, |
| maximum=121, |
| step=1, |
| value=DEFAULT_NUM_FRAMES, |
| label="Frames", |
| ) |
| |
| seed = gr.Number( |
| label="Seed", |
| value=-1, |
| precision=0, |
| ) |
| |
| resolution = gr.Dropdown( |
| label="Resolution", |
| choices=VIDEO_RESOLUTION_CHOICES, |
| value=DEFAULT_RESOLUTION, |
| ) |
| |
| validation_num_timesteps = gr.Slider( |
| minimum=1, |
| maximum=100, |
| step=1, |
| value=DEFAULT_TIMESTEPS, |
| label="Timesteps", |
| ) |
| |
| validation_timestep_shift = gr.Number( |
| label="Timestep Shift", |
| value=DEFAULT_TIMESTEP_SHIFT, |
| ) |
| |
| cfg_text_scale = gr.Number( |
| label="CFG Text Scale", |
| value=DEFAULT_CFG_TEXT_SCALE, |
| ) |
| |
| run_button = gr.Button( |
| "Run", |
| variant="primary", |
| ) |
| |
| with gr.Column(): |
| |
| output_video = gr.Video( |
| label="Generated Video" |
| ) |
| |
| output_image = gr.Image( |
| label="Generated Image" |
| ) |
| |
| output_text = gr.Textbox( |
| label="Text Output", |
| lines=8, |
| ) |
| |
| status = gr.Markdown() |
| |
| logs = gr.Textbox( |
| label="Logs", |
| lines=15, |
| ) |
| |
| run_button.click( |
| fn=run_task, |
| inputs=[ |
| task, |
| prompt, |
| input_image, |
| input_video, |
| question, |
| height, |
| width, |
| num_frames, |
| seed, |
| resolution, |
| validation_num_timesteps, |
| validation_timestep_shift, |
| cfg_text_scale, |
| ], |
| outputs=[ |
| output_video, |
| output_image, |
| output_text, |
| status, |
| logs, |
| ], |
| ) |
| |
| # ========================================================= |
| # LAUNCH |
| # ========================================================= |
| |
| demo.queue( |
| max_size=4, |
| default_concurrency_limit=1, |
| ).launch( |
| server_name="0.0.0.0", |
| server_port=7860, |
| ) |
