Spaces:
Running on Zero
Running on Zero
| import types | |
| import random | |
| import spaces | |
| import logging | |
| import os | |
| from pathlib import Path | |
| from datetime import datetime | |
| import re | |
| import torch | |
| import numpy as np | |
| from diffusers import AutoencoderKLWan, UniPCMultistepScheduler | |
| from diffusers.utils import export_to_video | |
| from diffusers import AutoModel | |
| import gradio as gr | |
| import json | |
| import tempfile | |
| from huggingface_hub import hf_hub_download | |
| import traceback | |
| import torch.nn.functional as F | |
| original_sdpa = F.scaled_dot_product_attention | |
| def patched_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=None): | |
| kwargs = {} | |
| if attn_mask is not None: | |
| kwargs['attn_mask'] = attn_mask | |
| if dropout_p != 0.0: | |
| kwargs['dropout_p'] = dropout_p | |
| if is_causal: | |
| kwargs['is_causal'] = is_causal | |
| if scale is not None: | |
| kwargs['scale'] = scale | |
| return original_sdpa(query, key, value, **kwargs) | |
| F.scaled_dot_product_attention = patched_scaled_dot_product_attention | |
| from src.pipeline_wan_nag import NAGWanPipeline | |
| from src.transformer_wan_nag import NagWanTransformer3DModel | |
| MOD_VALUE = 32 | |
| DEFAULT_DURATION_SECONDS = 7.9 | |
| DEFAULT_STEPS = 7 | |
| DEFAULT_SEED = 0 | |
| DEFAULT_H_SLIDER_VALUE = 896 | |
| DEFAULT_W_SLIDER_VALUE = 384 | |
| SLIDER_MIN_H, SLIDER_MAX_H = 128, 1280 | |
| SLIDER_MIN_W, SLIDER_MAX_W = 128, 1280 | |
| MAX_SEED = np.iinfo(np.int32).max | |
| FIXED_FPS = 16 | |
| MIN_FRAMES_MODEL = 8 | |
| MAX_FRAMES_MODEL = 240 | |
| MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers" | |
| SUB_MODEL_ID = "vrgamedevgirl84/Wan14BT2VFusioniX" | |
| SUB_MODEL_FILENAME = "Wan14BT2VFusioniX_fp16_.safetensors" | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| log = logging.getLogger() | |
| device = 'cuda' | |
| dtype = torch.bfloat16 | |
| try: | |
| vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16) | |
| wan_path = hf_hub_download(repo_id=SUB_MODEL_ID, filename=SUB_MODEL_FILENAME) | |
| transformer = NagWanTransformer3DModel.from_single_file(wan_path, torch_dtype=torch.bfloat16) | |
| pipe = NAGWanPipeline.from_pretrained(MODEL_ID, vae=vae, transformer=transformer, torch_dtype=torch.bfloat16) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0) | |
| pipe.to("cuda") | |
| pipe.transformer.__class__.attn_processors = NagWanTransformer3DModel.attn_processors | |
| pipe.transformer.__class__.set_attn_processor = NagWanTransformer3DModel.set_attn_processor | |
| pipe.transformer.__class__.forward = NagWanTransformer3DModel.forward | |
| print("NAG Video Model loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading NAG Video Model: {e}") | |
| pipe = None | |
| def get_duration(prompt, nag_negative_prompt, nag_scale, height, width, duration_seconds, GPU_time, steps, seed, randomize_seed): | |
| GPU_time = float(GPU_time) | |
| if GPU_time == 0: | |
| video_duration = int(duration_seconds) * int(steps) * 2.25 + 5 | |
| else: | |
| video_duration = GPU_time / 1.5 | |
| gr.Info(f"GPU time = {video_duration * 1.5}s") | |
| return video_duration | |
| def generate_video( | |
| prompt, | |
| nag_negative_prompt, | |
| nag_scale, | |
| height=DEFAULT_H_SLIDER_VALUE, | |
| width=DEFAULT_W_SLIDER_VALUE, | |
| duration_seconds=DEFAULT_DURATION_SECONDS, | |
| GPU_time=60, | |
| steps=DEFAULT_STEPS, | |
| seed=DEFAULT_SEED, | |
| randomize_seed=False, | |
| ): | |
| if pipe is None: | |
| return None, DEFAULT_SEED | |
| try: | |
| target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE) | |
| target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE) | |
| num_frames = np.clip(int(round(int(duration_seconds) * FIXED_FPS) + 1), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL) | |
| current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) | |
| with torch.inference_mode(): | |
| nag_output_frames_list = pipe( | |
| prompt=prompt, | |
| nag_negative_prompt=nag_negative_prompt, | |
| nag_scale=nag_scale, | |
| nag_tau=3.5, | |
| nag_alpha=0.5, | |
| height=target_h, width=target_w, num_frames=num_frames, | |
| guidance_scale=0., | |
| num_inference_steps=int(steps), | |
| generator=torch.Generator(device="cuda").manual_seed(current_seed) | |
| ).frames[0] | |
| video_filename = f"{current_seed}_{steps}_{nag_scale}.mp4" | |
| video_path = os.path.join(tempfile.gettempdir(), video_filename) | |
| export_to_video(nag_output_frames_list, video_path, fps=FIXED_FPS, quality=8) | |
| return video_path, current_seed | |
| except Exception as e: | |
| print(f"Error in video generation: {e}") | |
| return None, current_seed | |
| with open("theme.json", "r") as f: | |
| theme_data = json.load(f) | |
| theme = gr.themes.Base.from_dict(theme_data) | |
| with gr.Blocks(theme=theme) as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Prompt",lines=2,elem_classes="prompt-input") | |
| duration_seconds_input = gr.Slider(minimum=0.5,maximum=15,step=0.1,value=DEFAULT_DURATION_SECONDS,label="Duration",elem_classes="slider-container") | |
| GPU_time_input = gr.Slider(value=171.0,minimum=0.0,maximum=300.0,step=1.0,label="GPU time (s)") | |
| randomize_seed_checkbox = gr.Checkbox(label="Random Seed",value=True,interactive=True) | |
| seed_input = gr.Slider(label="Seed",minimum=0,maximum=MAX_SEED,step=1,value=DEFAULT_SEED,interactive=True) | |
| generate_button = gr.Button("Generate Video",variant="primary",elem_classes="generate-btn") | |
| with gr.Accordion("Advanced Settings", open=True): | |
| nag_negative_prompt = gr.Textbox(label="Negative Prompt",value="",lines=2,) | |
| height_input = gr.Slider(minimum=SLIDER_MIN_H,maximum=SLIDER_MAX_H,step=MOD_VALUE,value=DEFAULT_H_SLIDER_VALUE,label=f"Height",elem_classes="slider-container") | |
| width_input = gr.Slider(minimum=SLIDER_MIN_W,maximum=SLIDER_MAX_W,step=MOD_VALUE,value=DEFAULT_W_SLIDER_VALUE,label=f"Width",elem_classes="slider-container") | |
| steps_slider = gr.Slider(minimum=1,maximum=10,step=1,value=DEFAULT_STEPS,label="Inference Steps",elem_classes="slider-container") | |
| nag_scale = gr.Slider(label="NAG Scale",minimum=1.0,maximum=25.0,step=1,value=18.0,) | |
| with gr.Column(): | |
| video_output = gr.Video(label="", autoplay=False, interactive=False, loop=True, width=480) | |
| ui_inputs = [ | |
| prompt, nag_negative_prompt, nag_scale, height_input, width_input, | |
| duration_seconds_input, GPU_time_input, steps_slider,seed_input, randomize_seed_checkbox, | |
| ] | |
| generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input], show_progress="full",) | |
| if __name__ == "__main__": | |
| demo.queue().launch() |