Ani14/Video-agent
0
1"""2Model handler for WAN-VACE video generation3"""4import torch5 6# -----------------------------------------------------------------------------7# XPU shim for CPU‑only environments8#9# Newer versions of `diffusers` attempt to call `torch.xpu.empty_cache()` for10# Intel GPU support. If the installed PyTorch build does not include XPU11# support (as is the case on CPU‑only environments), accessing `torch.xpu`12# results in an AttributeError. To avoid this, we define a dummy `xpu`13# namespace on the `torch` module when it is missing. This namespace14# implements the minimal methods used by `diffusers` (`empty_cache`,15# `is_available`, and `device_count`).16#17# Intel’s `intel-extension-for-pytorch` provides XPU support, but even when18# installed, some CPU builds of PyTorch may not expose `torch.xpu`. This19# shim ensures that the application runs regardless of whether XPU support is20# present.21# -----------------------------------------------------------------------------22if not hasattr(torch, "xpu"):23 class _DummyXPU:24 @staticmethod25 def empty_cache():26 return None27 @staticmethod28 def manual_seed(_seed: int):29 return None30 @staticmethod31 def is_available():32 return False33 @staticmethod34 def device_count():35 return 036 @staticmethod37 def current_device():38 return 039 @staticmethod40 def set_device(_idx: int):41 return None42 torch.xpu = _DummyXPU() # type: ignore43import time44from typing import Optional, Tuple, Any45from transformers import UMT5EncoderModel46from diffusers import AutoencoderKLWan, WanVACEPipeline, WanVACETransformer3DModel, GGUFQuantizationConfig47from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler48from diffusers.utils import export_to_video49from huggingface_hub import login50import gradio as gr51 52from config import MODEL_CONFIG, DEFAULT_PARAMS, HF_TOKEN53import os54from utils import create_temp_video_path, validate_generation_params, validate_prompt, format_generation_info55 56class WanVACEModelHandler:57 """Handler for WAN-VACE model loading and video generation"""58 59 def __init__(self):60 self.pipe = None61 self.is_loaded = False62 self.loading_progress = 063 64 def login_hf(self) -> bool:65 """Login to Hugging Face"""66 try:67 login(token=HF_TOKEN)68 return True69 except Exception as e:70 print(f"Warning: Could not login to Hugging Face: {e}")71 return False72 73 def load_model(self, progress_callback=None) -> Tuple[bool, str]:74 """Load the WAN-VACE model components"""75 try:76 # Login to HF77 self.login_hf()78 79 if progress_callback:80 progress_callback(0.1, "Loading transformer model...")81 82 # Determine desired dtype for CPU/GPU execution. 83 # Hugging Face Spaces often run on CPU, where bfloat16 may not be supported. 84 # Allow the dtype to be configured via the WAN_DTYPE environment variable. 85 # Supported values: "bfloat16" (default) or "float32". 86 dtype_str = os.getenv("WAN_DTYPE", "bfloat16").lower()87 # Select compute dtype: use bfloat16 only if requested and available. 88 # Fall back to float32 otherwise. 89 compute_dtype = torch.bfloat16 if dtype_str == "bfloat16" else torch.float3290 # Likewise for the torch dtype used when loading weights. 91 torch_dtype = compute_dtype92 93 # Load transformer94 transformer = WanVACETransformer3DModel.from_single_file(95 MODEL_CONFIG["transformer_path"],96 quantization_config=GGUFQuantizationConfig(compute_dtype=compute_dtype),97 torch_dtype=torch_dtype,98 )99 100 if progress_callback:101 progress_callback(0.4, "Loading text encoder...")102 103 # Load text encoder104 text_encoder = UMT5EncoderModel.from_pretrained(105 MODEL_CONFIG["text_encoder_path"],106 gguf_file=MODEL_CONFIG["text_encoder_file"],107 torch_dtype=torch_dtype,108 )109 110 if progress_callback:111 progress_callback(0.7, "Loading VAE...")112 113 # Load VAE114 vae = AutoencoderKLWan.from_pretrained(115 MODEL_CONFIG["vae_path"],116 subfolder="vae",117 torch_dtype=torch.float32118 )119 120 if progress_callback:121 progress_callback(0.9, "Assembling pipeline...")122 123 # Create pipeline124 self.pipe = WanVACEPipeline.from_pretrained(125 MODEL_CONFIG["pipeline_path"],126 transformer=transformer,127 text_encoder=text_encoder,128 vae=vae,129 torch_dtype=torch_dtype130 )131 132 # Configure scheduler133 flow_shift = DEFAULT_PARAMS["flow_shift"]134 self.pipe.scheduler = UniPCMultistepScheduler.from_config(135 self.pipe.scheduler.config, 136 flow_shift=flow_shift137 )138 139 # Enable optimizations140 self.pipe.enable_model_cpu_offload()141 self.pipe.vae.enable_tiling()142 143 self.is_loaded = True144 145 if progress_callback:146 progress_callback(1.0, "Model loaded successfully!")147 148 return True, "Model loaded successfully!"149 150 except Exception as e:151 error_msg = f"Error loading model: {str(e)}"152 if progress_callback:153 progress_callback(0, error_msg)154 return False, error_msg155 156 def generate_video(157 self,158 prompt: str,159 negative_prompt: str = "",160 width: int = DEFAULT_PARAMS["width"],161 height: int = DEFAULT_PARAMS["height"],162 num_frames: int = DEFAULT_PARAMS["num_frames"],163 num_inference_steps: int = DEFAULT_PARAMS["num_inference_steps"],164 guidance_scale: float = DEFAULT_PARAMS["guidance_scale"],165 seed: Optional[int] = None,166 progress_callback=None167 ) -> Tuple[bool, str, str, str]:168 """169 Generate video from text prompt170 Returns: (success, video_path, error_message, generation_info)171 """172 173 if not self.is_loaded:174 return False, "", "Model not loaded. Please load the model first.", ""175 176 # Validate inputs177 prompt_valid, prompt_error = validate_prompt(prompt)178 if not prompt_valid:179 return False, "", prompt_error or "Invalid prompt", ""180 181 params_valid, params_error = validate_generation_params(182 width, height, num_frames, num_inference_steps, guidance_scale183 )184 if not params_valid:185 return False, "", params_error or "Invalid parameters", ""186 187 try:188 if progress_callback:189 progress_callback(0.1, "Preparing generation...")190 191 # Check if pipeline is loaded192 if self.pipe is None:193 return False, "", "Pipeline not initialized. Please load the model first.", ""194 195 # Set up generator with seed196 generator = torch.Generator()197 if seed is not None:198 generator.manual_seed(seed)199 else:200 generator.manual_seed(0) # Default seed201 202 if progress_callback:203 progress_callback(0.2, "Starting video generation...")204 205 start_time = time.time()206 207 # Generate video208 output = self.pipe(209 prompt=prompt,210 negative_prompt=negative_prompt if negative_prompt else None,211 width=width,212 height=height,213 num_frames=num_frames,214 num_inference_steps=num_inference_steps,215 guidance_scale=guidance_scale,216 conditioning_scale=DEFAULT_PARAMS["conditioning_scale"],217 generator=generator,218 ).frames[0]219 220 if progress_callback:221 progress_callback(0.8, "Exporting video...")222 223 # Export to video file224 output_path = create_temp_video_path()225 export_to_video(output, output_path, fps=DEFAULT_PARAMS["fps"])226 227 generation_time = time.time() - start_time228 229 if progress_callback:230 progress_callback(1.0, "Video generation complete!")231 232 # Format generation info233 gen_info = format_generation_info(234 prompt, negative_prompt, width, height, num_frames,235 num_inference_steps, guidance_scale, generation_time236 )237 238 return True, output_path, "", gen_info239 240 except Exception as e:241 error_msg = f"Error during video generation: {str(e)}"242 if progress_callback:243 progress_callback(0, error_msg)244 return False, "", error_msg, ""245 246# Global model handler instance247model_handler = WanVACEModelHandler()248