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 @spaces.GPU(duration=get_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()