ALSv/self-forcing
0
1"""2Demo for Self-Forcing.3"""4 5import os6import re7import random8import time9import base6410import argparse11import hashlib12import subprocess13import urllib.request14from io import BytesIO15from PIL import Image16import numpy as np17import torch18from omegaconf import OmegaConf19from flask import Flask, render_template, jsonify20from flask_socketio import SocketIO, emit21import queue22from threading import Thread, Event23 24from pipeline import CausalInferencePipeline25from demo_utils.constant import ZERO_VAE_CACHE26from demo_utils.vae_block3 import VAEDecoderWrapper27from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder28from demo_utils.utils import generate_timestamp29from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation30 31# Parse arguments32parser = argparse.ArgumentParser()33parser.add_argument('--port', type=int, default=5001)34parser.add_argument('--host', type=str, default='0.0.0.0')35parser.add_argument("--checkpoint_path", type=str, default='./checkpoints/self_forcing_dmd.pt')36parser.add_argument("--config_path", type=str, default='./configs/self_forcing_dmd.yaml')37parser.add_argument('--trt', action='store_true')38args = parser.parse_args()39 40print(f'Free VRAM {get_cuda_free_memory_gb(gpu)} GB')41low_memory = get_cuda_free_memory_gb(gpu) < 4042 43# Load models44config = OmegaConf.load(args.config_path)45default_config = OmegaConf.load("configs/default_config.yaml")46config = OmegaConf.merge(default_config, config)47 48text_encoder = WanTextEncoder()49 50# Global variables for dynamic model switching51current_vae_decoder = None52current_use_taehv = False53fp8_applied = False54torch_compile_applied = False55global frame_number56frame_number = 057anim_name = ""58frame_rate = 659 60def initialize_vae_decoder(use_taehv=False, use_trt=False):61 """Initialize VAE decoder based on the selected option"""62 global current_vae_decoder, current_use_taehv63 64 if use_trt:65 from demo_utils.vae import VAETRTWrapper66 current_vae_decoder = VAETRTWrapper()67 return current_vae_decoder68 69 if use_taehv:70 from demo_utils.taehv import TAEHV71 # Check if taew2_1.pth exists in checkpoints folder, download if missing72 taehv_checkpoint_path = "checkpoints/taew2_1.pth"73 if not os.path.exists(taehv_checkpoint_path):74 print(f"taew2_1.pth not found in checkpoints folder {taehv_checkpoint_path}. Downloading...")75 os.makedirs("checkpoints", exist_ok=True)76 download_url = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth"77 try:78 urllib.request.urlretrieve(download_url, taehv_checkpoint_path)79 print(f"Successfully downloaded taew2_1.pth to {taehv_checkpoint_path}")80 except Exception as e:81 print(f"Failed to download taew2_1.pth: {e}")82 raise83 84 class DotDict(dict):85 __getattr__ = dict.__getitem__86 __setattr__ = dict.__setitem__87 88 class TAEHVDiffusersWrapper(torch.nn.Module):89 def __init__(self):90 super().__init__()91 self.dtype = torch.float1692 self.taehv = TAEHV(checkpoint_path=taehv_checkpoint_path).to(self.dtype)93 self.config = DotDict(scaling_factor=1.0)94 95 def decode(self, latents, return_dict=None):96 # n, c, t, h, w = latents.shape97 # low-memory, set parallel=True for faster + higher memory98 return self.taehv.decode_video(latents, parallel=False).mul_(2).sub_(1)99 100 current_vae_decoder = TAEHVDiffusersWrapper()101 else:102 current_vae_decoder = VAEDecoderWrapper()103 vae_state_dict = torch.load('wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', map_location="cpu")104 decoder_state_dict = {}105 for key, value in vae_state_dict.items():106 if 'decoder.' in key or 'conv2' in key:107 decoder_state_dict[key] = value108 current_vae_decoder.load_state_dict(decoder_state_dict)109 110 current_vae_decoder.eval()111 current_vae_decoder.to(dtype=torch.float16)112 current_vae_decoder.requires_grad_(False)113 current_vae_decoder.to(gpu)114 current_use_taehv = use_taehv115 116 print(f"✅ VAE decoder initialized with {'TAEHV' if use_taehv else 'default VAE'}")117 return current_vae_decoder118 119 120# Initialize with default VAE121vae_decoder = initialize_vae_decoder(use_taehv=False, use_trt=args.trt)122 123transformer = WanDiffusionWrapper(is_causal=True)124state_dict = torch.load(args.checkpoint_path, map_location="cpu")125transformer.load_state_dict(state_dict['generator_ema'])126 127text_encoder.eval()128transformer.eval()129 130transformer.to(dtype=torch.float16)131text_encoder.to(dtype=torch.bfloat16)132 133text_encoder.requires_grad_(False)134transformer.requires_grad_(False)135 136pipeline = CausalInferencePipeline(137 config,138 device=gpu,139 generator=transformer,140 text_encoder=text_encoder,141 vae=vae_decoder142)143 144if low_memory:145 DynamicSwapInstaller.install_model(text_encoder, device=gpu)146else:147 text_encoder.to(gpu)148transformer.to(gpu)149 150# Flask and SocketIO setup151app = Flask(__name__)152app.config['SECRET_KEY'] = 'frontend_buffered_demo'153socketio = SocketIO(app, cors_allowed_origins="*")154 155generation_active = False156stop_event = Event()157frame_send_queue = queue.Queue()158sender_thread = None159models_compiled = False160 161 162def tensor_to_base64_frame(frame_tensor):163 """Convert a single frame tensor to base64 image string."""164 global frame_number, anim_name165 # Clamp and normalize to 0-255166 frame = torch.clamp(frame_tensor.float(), -1., 1.) * 127.5 + 127.5167 frame = frame.to(torch.uint8).cpu().numpy()168 169 # CHW -> HWC170 if len(frame.shape) == 3:171 frame = np.transpose(frame, (1, 2, 0))172 173 # Convert to PIL Image174 if frame.shape[2] == 3: # RGB175 image = Image.fromarray(frame, 'RGB')176 else: # Handle other formats177 image = Image.fromarray(frame)178 179 # Convert to base64180 buffer = BytesIO()181 image.save(buffer, format='JPEG', quality=100)182 if not os.path.exists("./images/%s" % anim_name):183 os.makedirs("./images/%s" % anim_name)184 frame_number += 1185 image.save("./images/%s/%s_%03d.jpg" % (anim_name, anim_name, frame_number))186 img_str = base64.b64encode(buffer.getvalue()).decode()187 return f"data:image/jpeg;base64,{img_str}"188 189 190def frame_sender_worker():191 """Background thread that processes frame send queue non-blocking."""192 global frame_send_queue, generation_active, stop_event193 194 print("📡 Frame sender thread started")195 196 while True:197 frame_data = None198 try:199 # Get frame data from queue200 frame_data = frame_send_queue.get(timeout=1.0)201 202 if frame_data is None: # Shutdown signal203 frame_send_queue.task_done() # Mark shutdown signal as done204 break205 206 frame_tensor, frame_index, block_index, job_id = frame_data207 208 # Convert tensor to base64209 base64_frame = tensor_to_base64_frame(frame_tensor)210 211 # Send via SocketIO212 try:213 socketio.emit('frame_ready', {214 'data': base64_frame,215 'frame_index': frame_index,216 'block_index': block_index,217 'job_id': job_id218 })219 except Exception as e:220 print(f"⚠️ Failed to send frame {frame_index}: {e}")221 222 frame_send_queue.task_done()223 224 except queue.Empty:225 # Check if we should continue running226 if not generation_active and frame_send_queue.empty():227 break228 except Exception as e:229 print(f"❌ Frame sender error: {e}")230 # Make sure to mark task as done even if there's an error231 if frame_data is not None:232 try:233 frame_send_queue.task_done()234 except Exception as e:235 print(f"❌ Failed to mark frame task as done: {e}")236 break237 238 print("📡 Frame sender thread stopped")239 240 241@torch.no_grad()242def generate_video_stream(prompt, seed, enable_torch_compile=False, enable_fp8=False, use_taehv=False):243 """Generate video and push frames immediately to frontend."""244 global generation_active, stop_event, frame_send_queue, sender_thread, models_compiled, torch_compile_applied, fp8_applied, current_vae_decoder, current_use_taehv, frame_rate, anim_name245 246 try:247 generation_active = True248 stop_event.clear()249 job_id = generate_timestamp()250 251 # Start frame sender thread if not already running252 if sender_thread is None or not sender_thread.is_alive():253 sender_thread = Thread(target=frame_sender_worker, daemon=True)254 sender_thread.start()255 256 # Emit progress updates257 def emit_progress(message, progress):258 try:259 socketio.emit('progress', {260 'message': message,261 'progress': progress,262 'job_id': job_id263 })264 except Exception as e:265 print(f"❌ Failed to emit progress: {e}")266 267 emit_progress('Starting generation...', 0)268 269 # Handle VAE decoder switching270 if use_taehv != current_use_taehv:271 emit_progress('Switching VAE decoder...', 2)272 print(f"🔄 Switching VAE decoder to {'TAEHV' if use_taehv else 'default VAE'}")273 current_vae_decoder = initialize_vae_decoder(use_taehv=use_taehv)274 # Update pipeline with new VAE decoder275 pipeline.vae = current_vae_decoder276 277 # Handle FP8 quantization278 if enable_fp8 and not fp8_applied:279 emit_progress('Applying FP8 quantization...', 3)280 print("🔧 Applying FP8 quantization to transformer")281 from torchao.quantization.quant_api import quantize_, Float8DynamicActivationFloat8WeightConfig, PerTensor282 quantize_(transformer, Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()))283 fp8_applied = True284 285 # Text encoding286 emit_progress('Encoding text prompt...', 8)287 conditional_dict = text_encoder(text_prompts=[prompt])288 for key, value in conditional_dict.items():289 conditional_dict[key] = value.to(dtype=torch.float16)290 if low_memory:291 gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5292 move_model_to_device_with_memory_preservation(293 text_encoder,target_device=gpu, preserved_memory_gb=gpu_memory_preservation)294 295 # Handle torch.compile if enabled296 torch_compile_applied = enable_torch_compile297 if enable_torch_compile and not models_compiled:298 # Compile transformer and decoder299 transformer.compile(mode="max-autotune-no-cudagraphs")300 if not current_use_taehv and not low_memory and not args.trt:301 current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")302 303 # Initialize generation304 emit_progress('Initializing generation...', 12)305 306 rnd = torch.Generator(gpu).manual_seed(seed)307 # all_latents = torch.zeros([1, 21, 16, 60, 104], device=gpu, dtype=torch.bfloat16)308 309 pipeline._initialize_kv_cache(batch_size=1, dtype=torch.float16, device=gpu)310 pipeline._initialize_crossattn_cache(batch_size=1, dtype=torch.float16, device=gpu)311 312 noise = torch.randn([1, 21, 16, 60, 104], device=gpu, dtype=torch.float16, generator=rnd)313 314 # Generation parameters315 num_blocks = 7316 current_start_frame = 0317 num_input_frames = 0318 all_num_frames = [pipeline.num_frame_per_block] * num_blocks319 if current_use_taehv:320 vae_cache = None321 else:322 vae_cache = ZERO_VAE_CACHE323 for i in range(len(vae_cache)):324 vae_cache[i] = vae_cache[i].to(device=gpu, dtype=torch.float16)325 326 total_frames_sent = 0327 generation_start_time = time.time()328 329 emit_progress('Generating frames... (frontend handles timing)', 15)330 331 for idx, current_num_frames in enumerate(all_num_frames):332 if not generation_active or stop_event.is_set():333 break334 335 progress = int(((idx + 1) / len(all_num_frames)) * 80) + 15336 337 # Special message for first block with torch.compile338 if idx == 0 and torch_compile_applied and not models_compiled:339 emit_progress(340 f'Processing block 1/{len(all_num_frames)} - Compiling models (may take 5-10 minutes)...', progress)341 print(f"🔥 Processing block {idx+1}/{len(all_num_frames)}")342 models_compiled = True343 else:344 emit_progress(f'Processing block {idx+1}/{len(all_num_frames)}...', progress)345 print(f"🔄 Processing block {idx+1}/{len(all_num_frames)}")346 347 block_start_time = time.time()348 349 noisy_input = noise[:, current_start_frame -350 num_input_frames:current_start_frame + current_num_frames - num_input_frames]351 352 # Denoising loop353 denoising_start = time.time()354 for index, current_timestep in enumerate(pipeline.denoising_step_list):355 if not generation_active or stop_event.is_set():356 break357 358 timestep = torch.ones([1, current_num_frames], device=noise.device,359 dtype=torch.int64) * current_timestep360 361 if index < len(pipeline.denoising_step_list) - 1:362 _, denoised_pred = transformer(363 noisy_image_or_video=noisy_input,364 conditional_dict=conditional_dict,365 timestep=timestep,366 kv_cache=pipeline.kv_cache1,367 crossattn_cache=pipeline.crossattn_cache,368 current_start=current_start_frame * pipeline.frame_seq_length369 )370 next_timestep = pipeline.denoising_step_list[index + 1]371 noisy_input = pipeline.scheduler.add_noise(372 denoised_pred.flatten(0, 1),373 torch.randn_like(denoised_pred.flatten(0, 1)),374 next_timestep * torch.ones([1 * current_num_frames], device=noise.device, dtype=torch.long)375 ).unflatten(0, denoised_pred.shape[:2])376 else:377 _, denoised_pred = transformer(378 noisy_image_or_video=noisy_input,379 conditional_dict=conditional_dict,380 timestep=timestep,381 kv_cache=pipeline.kv_cache1,382 crossattn_cache=pipeline.crossattn_cache,383 current_start=current_start_frame * pipeline.frame_seq_length384 )385 386 if not generation_active or stop_event.is_set():387 break388 389 denoising_time = time.time() - denoising_start390 print(f"⚡ Block {idx+1} denoising completed in {denoising_time:.2f}s")391 392 # Record output393 # all_latents[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred394 395 # Update KV cache for next block396 if idx != len(all_num_frames) - 1:397 transformer(398 noisy_image_or_video=denoised_pred,399 conditional_dict=conditional_dict,400 timestep=torch.zeros_like(timestep),401 kv_cache=pipeline.kv_cache1,402 crossattn_cache=pipeline.crossattn_cache,403 current_start=current_start_frame * pipeline.frame_seq_length,404 )405 406 # Decode to pixels and send frames immediately407 print(f"🎨 Decoding block {idx+1} to pixels...")408 decode_start = time.time()409 if args.trt:410 all_current_pixels = []411 for i in range(denoised_pred.shape[1]):412 is_first_frame = torch.tensor(1.0).cuda().half() if idx == 0 and i == 0 else \413 torch.tensor(0.0).cuda().half()414 outputs = vae_decoder.forward(denoised_pred[:, i:i + 1, :, :, :].half(), is_first_frame, *vae_cache)415 # outputs = vae_decoder.forward(denoised_pred.float(), *vae_cache)416 current_pixels, vae_cache = outputs[0], outputs[1:]417 print(current_pixels.max(), current_pixels.min())418 all_current_pixels.append(current_pixels.clone())419 pixels = torch.cat(all_current_pixels, dim=1)420 if idx == 0:421 pixels = pixels[:, 3:, :, :, :] # Skip first 3 frames of first block422 else:423 if current_use_taehv:424 if vae_cache is None:425 vae_cache = denoised_pred426 else:427 denoised_pred = torch.cat([vae_cache, denoised_pred], dim=1)428 vae_cache = denoised_pred[:, -3:, :, :, :]429 pixels = current_vae_decoder.decode(denoised_pred)430 print(f"denoised_pred shape: {denoised_pred.shape}")431 print(f"pixels shape: {pixels.shape}")432 if idx == 0:433 pixels = pixels[:, 3:, :, :, :] # Skip first 3 frames of first block434 else:435 pixels = pixels[:, 12:, :, :, :]436 437 else:438 pixels, vae_cache = current_vae_decoder(denoised_pred.half(), *vae_cache)439 if idx == 0:440 pixels = pixels[:, 3:, :, :, :] # Skip first 3 frames of first block441 442 decode_time = time.time() - decode_start443 print(f"🎨 Block {idx+1} VAE decoding completed in {decode_time:.2f}s")444 445 # Queue frames for non-blocking sending446 block_frames = pixels.shape[1]447 print(f"📡 Queueing {block_frames} frames from block {idx+1} for sending...")448 queue_start = time.time()449 450 for frame_idx in range(block_frames):451 if not generation_active or stop_event.is_set():452 break453 454 frame_tensor = pixels[0, frame_idx].cpu()455 456 # Queue frame data in non-blocking way457 frame_send_queue.put((frame_tensor, total_frames_sent, idx, job_id))458 total_frames_sent += 1459 460 queue_time = time.time() - queue_start461 block_time = time.time() - block_start_time462 print(f"✅ Block {idx+1} completed in {block_time:.2f}s ({block_frames} frames queued in {queue_time:.3f}s)")463 464 current_start_frame += current_num_frames465 466 generation_time = time.time() - generation_start_time467 print(f"🎉 Generation completed in {generation_time:.2f}s! {total_frames_sent} frames queued for sending")468 469 # Wait for all frames to be sent before completing470 emit_progress('Waiting for all frames to be sent...', 97)471 print("⏳ Waiting for all frames to be sent...")472 frame_send_queue.join() # Wait for all queued frames to be processed473 print("✅ All frames sent successfully!")474 475 generate_mp4_from_images("./images","./videos/"+anim_name+".mp4", frame_rate )476 # Final progress update477 emit_progress('Generation complete!', 100)478 479 try:480 socketio.emit('generation_complete', {481 'message': 'Video generation completed!',482 'total_frames': total_frames_sent,483 'generation_time': f"{generation_time:.2f}s",484 'job_id': job_id485 })486 except Exception as e:487 print(f"❌ Failed to emit generation complete: {e}")488 489 except Exception as e:490 print(f"❌ Generation failed: {e}")491 try:492 socketio.emit('error', {493 'message': f'Generation failed: {str(e)}',494 'job_id': job_id495 })496 except Exception as e:497 print(f"❌ Failed to emit error: {e}")498 finally:499 generation_active = False500 stop_event.set()501 502 # Clean up sender thread503 try:504 frame_send_queue.put(None)505 except Exception as e:506 print(f"❌ Failed to put None in frame_send_queue: {e}")507 508 509def generate_mp4_from_images(image_directory, output_video_path, fps=24):510 """511 Generate an MP4 video from a directory of images ordered alphabetically.512 513 :param image_directory: Path to the directory containing images.514 :param output_video_path: Path where the output MP4 will be saved.515 :param fps: Frames per second for the output video.516 """517 global anim_name518 # Construct the ffmpeg command519 cmd = [520 'ffmpeg',521 '-framerate', str(fps),522 '-i', os.path.join(image_directory, anim_name+'/'+anim_name+'_%03d.jpg'), # Adjust the pattern if necessary523 '-c:v', 'libx264',524 '-pix_fmt', 'yuv420p',525 output_video_path526 ]527 try:528 subprocess.run(cmd, check=True)529 print(f"Video saved to {output_video_path}")530 except subprocess.CalledProcessError as e:531 print(f"An error occurred: {e}")532 533def calculate_sha256(data):534 # Convert data to bytes if it's not already535 if isinstance(data, str):536 data = data.encode()537 # Calculate SHA-256 hash538 sha256_hash = hashlib.sha256(data).hexdigest()539 return sha256_hash540 541# Socket.IO event handlers542@socketio.on('connect')543def handle_connect():544 print('Client connected')545 emit('status', {'message': 'Connected to frontend-buffered demo server'})546 547 548@socketio.on('disconnect')549def handle_disconnect():550 print('Client disconnected')551 552 553@socketio.on('start_generation')554def handle_start_generation(data):555 global generation_active, frame_number, anim_name, frame_rate556 557 frame_number = 0558 if generation_active:559 emit('error', {'message': 'Generation already in progress'})560 return561 562 prompt = data.get('prompt', '')563 564 seed = data.get('seed', -1)565 if seed==-1:566 seed = random.randint(0, 2**32)567 568 # Extract words up to the first punctuation or newline569 words_up_to_punctuation = re.split(r'[^\w\s]', prompt)[0].strip() if prompt else ''570 if not words_up_to_punctuation:571 words_up_to_punctuation = re.split(r'[\n\r]', prompt)[0].strip()572 573 # Calculate SHA-256 hash of the entire prompt574 sha256_hash = calculate_sha256(prompt)575 576 # Create anim_name with the extracted words and first 10 characters of the hash577 anim_name = f"{words_up_to_punctuation[:20]}_{str(seed)}_{sha256_hash[:10]}"578 579 generation_active = True580 generation_start_time = time.time()581 enable_torch_compile = data.get('enable_torch_compile', False)582 enable_fp8 = data.get('enable_fp8', False)583 use_taehv = data.get('use_taehv', False)584 frame_rate = data.get('fps', 6)585 586 if not prompt:587 emit('error', {'message': 'Prompt is required'})588 return589 590 # Start generation in background thread591 socketio.start_background_task(generate_video_stream, prompt, seed,592 enable_torch_compile, enable_fp8, use_taehv)593 emit('status', {'message': 'Generation started - frames will be sent immediately'})594 595 596@socketio.on('stop_generation')597def handle_stop_generation():598 global generation_active, stop_event, frame_send_queue599 generation_active = False600 stop_event.set()601 602 # Signal sender thread to stop (will be processed after current frames)603 try:604 frame_send_queue.put(None)605 except Exception as e:606 print(f"❌ Failed to put None in frame_send_queue: {e}")607 608 emit('status', {'message': 'Generation stopped'})609 610# Web routes611 612 613@app.route('/')614def index():615 return render_template('demo.html')616 617 618@app.route('/api/status')619def api_status():620 return jsonify({621 'generation_active': generation_active,622 'free_vram_gb': get_cuda_free_memory_gb(gpu),623 'fp8_applied': fp8_applied,624 'torch_compile_applied': torch_compile_applied,625 'current_use_taehv': current_use_taehv626 })627 628 629if __name__ == '__main__':630 print(f"🚀 Starting demo on http://{args.host}:{args.port}")631 socketio.run(app, host=args.host, port=args.port, debug=False)632