Ani14/Video-agent
0
1"""2Gradio interface for WAN-VACE video generation3"""4import gradio as gr5import torch6 7# -----------------------------------------------------------------------------8# XPU shim for CPU‑only environments9#10# `diffusers` attempts to access `torch.xpu.empty_cache()` when cleaning up11# device memory. On CPU‑only builds of PyTorch (or builds without Intel12# extension support), the `xpu` attribute does not exist on the `torch`13# module. Defining a dummy `torch.xpu` prevents AttributeError during14# import.15# -----------------------------------------------------------------------------16if not hasattr(torch, "xpu"):17 class _DummyXPU:18 @staticmethod19 def empty_cache():20 return None21 @staticmethod22 def manual_seed(_seed: int):23 return None24 @staticmethod25 def is_available():26 return False27 @staticmethod28 def device_count():29 return 030 @staticmethod31 def current_device():32 return 033 @staticmethod34 def set_device(_idx: int):35 return None36 torch.xpu = _DummyXPU() # type: ignore37import time38from typing import Optional39 40# Import the simple planner41from planning import plan_from_topic42 43from config import UI_CONFIG, DEFAULT_PARAMS, SERVER_CONFIG44from model_handler import model_handler45from utils import cleanup_temp_files46 47def load_model_interface(progress=gr.Progress()):48 """Interface function for loading the model"""49 def progress_callback(value, message):50 progress(value, desc=message)51 52 success, message = model_handler.load_model(progress_callback)53 54 if success:55 return (56 gr.update(visible=False), # Hide load button57 gr.update(visible=True), # Show generation interface58 gr.update(value=message, visible=True), # Show success message59 gr.update(visible=False) # Hide error message60 )61 else:62 return (63 gr.update(visible=True), # Keep load button visible64 gr.update(visible=False), # Keep generation interface hidden65 gr.update(visible=False), # Hide success message66 gr.update(value=message, visible=True) # Show error message67 )68 69def generate_video_interface(70 prompt: str,71 negative_prompt: str,72 width: int,73 height: int,74 num_frames: int,75 num_inference_steps: int,76 guidance_scale: float,77 seed: Optional[int],78 progress=gr.Progress()79):80 """Interface function for video generation"""81 82 def progress_callback(value, message):83 progress(value, desc=message)84 85 # Plan the prompt: treat the user input as a high‑level concept and let the86 # agent craft a refined prompt and recommended negative prompt. If the user87 # supplies a negative prompt, it overrides the recommended negative prompt.88 plan = plan_from_topic(prompt)89 # Use the refined prompt from the plan90 effective_prompt = plan.prompt91 # If the user provided a negative prompt, use it; otherwise use the recommended one92 effective_negative = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else plan.negative_prompt93 94 success, video_path, error_msg, gen_info = model_handler.generate_video(95 prompt=effective_prompt,96 negative_prompt=effective_negative,97 width=width,98 height=height,99 num_frames=num_frames,100 num_inference_steps=num_inference_steps,101 guidance_scale=guidance_scale,102 seed=seed,103 progress_callback=progress_callback104 )105 106 if success:107 return (108 gr.update(value=video_path, visible=True), # Video output109 gr.update(value=gen_info, visible=True), # Generation info110 gr.update(visible=False) # Hide error message111 )112 else:113 return (114 gr.update(value=None, visible=False), # Hide video output115 gr.update(visible=False), # Hide generation info116 gr.update(value=error_msg, visible=True) # Show error message117 )118 119def create_interface():120 """Create the Gradio interface"""121 122 with gr.Blocks(123 title=UI_CONFIG["title"],124 theme=UI_CONFIG["theme"]125 ) as demo:126 127 # Header128 gr.Markdown(f"# {UI_CONFIG['title']}")129 gr.Markdown(UI_CONFIG["description"])130 131 # Model loading section132 with gr.Row():133 with gr.Column():134 load_btn = gr.Button(135 "🚀 Load Video Generation Model", 136 variant="primary", 137 size="lg"138 )139 load_success_msg = gr.Markdown(visible=False)140 load_error_msg = gr.Markdown(visible=False)141 142 # Main generation interface (initially hidden)143 with gr.Column(visible=False) as generation_interface:144 145 # Input section146 with gr.Row():147 with gr.Column(scale=2):148 with gr.Group():149 gr.Markdown("### 📝 Concept & Prompts")150 # The user supplies a high‑level concept or topic. The agent will151 # refine this into a detailed prompt automatically.152 prompt_input = gr.Textbox(153 label="Video Concept",154 placeholder="Describe the concept you want to generate, e.g. 'a pig in a winter forest'...",155 lines=3,156 value="a pig moving quickly in a beautiful winter scenery nature trees sunset tracking camera"157 )158 # Optional negative prompt: overrides the agent's recommended negative prompt.159 negative_prompt_input = gr.Textbox(160 label="Negative Prompt (Optional)",161 placeholder="Things you don't want in the video; leave empty to use the agent's recommendation...",162 lines=2,163 value=""164 )165 166 with gr.Column(scale=1):167 with gr.Group():168 gr.Markdown("### ⚙️ Generation Parameters")169 170 with gr.Row():171 width_slider = gr.Slider(172 label="Width",173 minimum=64,174 maximum=1920,175 step=8,176 value=DEFAULT_PARAMS["width"]177 )178 height_slider = gr.Slider(179 label="Height",180 minimum=64,181 maximum=1080,182 step=8,183 value=DEFAULT_PARAMS["height"]184 )185 186 num_frames_slider = gr.Slider(187 label="Number of Frames",188 minimum=1,189 maximum=200,190 step=1,191 value=DEFAULT_PARAMS["num_frames"]192 )193 194 inference_steps_slider = gr.Slider(195 label="Inference Steps",196 minimum=1,197 maximum=100,198 step=1,199 value=DEFAULT_PARAMS["num_inference_steps"]200 )201 202 guidance_scale_slider = gr.Slider(203 label="Guidance Scale",204 minimum=0.0,205 maximum=20.0,206 step=0.1,207 value=DEFAULT_PARAMS["guidance_scale"]208 )209 210 seed_input = gr.Number(211 label="Seed (Optional)",212 value=0,213 precision=0214 )215 216 # Generation button217 with gr.Row():218 generate_btn = gr.Button(219 "🎬 Generate Video",220 variant="primary",221 size="lg"222 )223 224 # Output section225 with gr.Row():226 with gr.Column():227 video_output = gr.Video(228 label="Generated Video",229 visible=False230 )231 232 generation_info = gr.Markdown(233 label="Generation Information",234 visible=False235 )236 237 generation_error = gr.Markdown(238 visible=False239 )240 241 # Additional controls242 with gr.Row():243 with gr.Column():244 gr.Markdown("""245 ### 💡 Tips:246 - Enter a short **concept** (e.g. “a busy city street at dawn”). The agent will expand it into a detailed prompt.247 - Adjust the **guidance scale**: higher values make the video adhere more closely to the refined prompt.248 - Increasing **inference steps** improves quality at the cost of generation time.249 - Use the optional **Negative Prompt** field only if you want to override the agent's recommended terms.250 - Keep width and height multiples of 8 for optimal performance.251 """)252 253 with gr.Column():254 if torch.cuda.is_available():255 gpu_info = f"🎮 GPU: {torch.cuda.get_device_name()}"256 else:257 gpu_info = "💻 Running on CPU"258 259 gr.Markdown(f"""260 ### 🖥️ System Information:261 {gpu_info}262 263 ### 📊 Model Information:264 - **Model:** WAN‑VACE 1.3B (Q4_0 Quantized)265 - **Text Encoder:** UMT5‑XXL266 - **Scheduler:** UniPC Multistep267 268 ### 🤖 Agent Details:269 - **Planning:** The agent automatically crafts a detailed prompt and a recommended negative prompt based on your concept.270 - **Override:** Supply your own negative prompt to override the recommendation if desired.271 """)272 273 # Event handlers274 load_btn.click(275 fn=load_model_interface,276 outputs=[277 load_btn,278 generation_interface,279 load_success_msg,280 load_error_msg281 ]282 )283 284 generate_btn.click(285 fn=generate_video_interface,286 inputs=[287 prompt_input,288 negative_prompt_input,289 width_slider,290 height_slider,291 num_frames_slider,292 inference_steps_slider,293 guidance_scale_slider,294 seed_input295 ],296 outputs=[297 video_output,298 generation_info,299 generation_error300 ]301 )302 303 return demo304 305def main():306 """Main function to launch the application"""307 print(f"🚀 Starting {UI_CONFIG['title']}...")308 print(f"🔧 Server configuration: {SERVER_CONFIG['host']}:{SERVER_CONFIG['port']}")309 310 # Check GPU availability311 if torch.cuda.is_available():312 print(f"🎮 GPU detected: {torch.cuda.get_device_name()}")313 print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")314 else:315 print("💻 Running on CPU (GPU recommended for better performance)")316 317 # Create interface and enable the event queue to support multiple users. 318 demo = create_interface()319 # Hugging Face Spaces expect `.queue()` to be called for handling request concurrency. 320 # Limiting concurrency_count to 1 helps prevent excessive memory usage on CPU-only hardware.321 demo = demo.queue()322 323 # Launch the interface. 324 demo.launch(325 server_name=SERVER_CONFIG["host"],326 server_port=SERVER_CONFIG["port"],327 share=SERVER_CONFIG["share"],328 show_error=True,329 )330 331if __name__ == "__main__":332 main()333