| |
|
| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from typing import List, Optional, Tuple, Dict, Any, Union |
| | import numpy as np |
| | from collections import deque |
| | import math |
| | from PIL import Image |
| | import torchvision.transforms as transforms |
| |
|
| | class TemporalConsistencyBuffer: |
| | """Enhanced temporal buffer for flexible batch generation""" |
| | |
| | def __init__(self, buffer_size: int = 8, feature_dim: int = 512): |
| | self.buffer_size = buffer_size |
| | self.feature_dim = feature_dim |
| | self.frame_features = deque(maxlen=buffer_size) |
| | self.frame_latents = deque(maxlen=buffer_size) |
| | self.frame_images = deque(maxlen=buffer_size) |
| | self.motion_vectors = deque(maxlen=buffer_size-1) |
| | self.temporal_weights = deque(maxlen=buffer_size) |
| | |
| | def add_frames(self, features: torch.Tensor, latents: torch.Tensor, images: Optional[torch.Tensor] = None, batch_size: int = 1): |
| | """Add batch of frames to temporal buffer""" |
| | for i in range(batch_size): |
| | frame_feat = features[i:i+1] if features.dim() > 3 else features |
| | frame_lat = latents[i:i+1] if latents.dim() > 3 else latents |
| | frame_img = images[i:i+1] if images is not None and images.dim() > 3 else images |
| | |
| | |
| | if len(self.frame_features) > 0: |
| | motion = frame_feat - self.frame_features[-1] |
| | self.motion_vectors.append(motion) |
| | |
| | self.frame_features.append(frame_feat) |
| | self.frame_latents.append(frame_lat) |
| | if frame_img is not None: |
| | self.frame_images.append(frame_img) |
| | |
| | |
| | weight = 1.0 / (len(self.frame_features) + 1) |
| | self.temporal_weights.append(weight) |
| | |
| | def get_reference_frame(self) -> Optional[torch.Tensor]: |
| | """Get the most recent frame as reference for I2V""" |
| | if len(self.frame_images) > 0: |
| | return self.frame_images[-1] |
| | elif len(self.frame_latents) > 0: |
| | return self.frame_latents[-1] |
| | return None |
| | |
| | def get_temporal_context(self, num_context_frames: int = 4) -> Dict[str, torch.Tensor]: |
| | """Get weighted temporal context for next frame batch""" |
| | if len(self.frame_features) == 0: |
| | return {"has_context": False} |
| | |
| | |
| | context_size = min(num_context_frames, len(self.frame_features)) |
| | recent_features = list(self.frame_features)[-context_size:] |
| | recent_latents = list(self.frame_latents)[-context_size:] |
| | recent_weights = list(self.temporal_weights)[-context_size:] |
| | |
| | |
| | stacked_features = torch.cat(recent_features, dim=0) |
| | stacked_latents = torch.cat(recent_latents, dim=0) |
| | weights = torch.tensor(recent_weights, device=stacked_features.device) |
| | |
| | |
| | predicted_motions = [] |
| | if len(self.motion_vectors) >= 2: |
| | |
| | recent_motions = list(self.motion_vectors)[-3:] |
| | for step in range(1, 4): |
| | if len(recent_motions) >= 2: |
| | |
| | motion_pred = ( |
| | recent_motions[-1] * 1.5 - |
| | recent_motions[-2] * 0.5 |
| | ) |
| | if len(recent_motions) >= 3: |
| | motion_pred += recent_motions[-3] * 0.1 |
| | else: |
| | motion_pred = recent_motions[-1] if recent_motions else None |
| | predicted_motions.append(motion_pred) |
| | |
| | return { |
| | "has_context": True, |
| | "frame_features": stacked_features, |
| | "frame_latents": stacked_latents, |
| | "temporal_weights": weights, |
| | "predicted_motions": predicted_motions, |
| | "sequence_length": len(self.frame_features), |
| | "reference_frame": self.get_reference_frame() |
| | } |
| |
|
| | class FlexibleTemporalAttention(nn.Module): |
| | """Flexible attention that handles variable batch sizes""" |
| | |
| | def __init__(self, dim: int, num_heads: int = 8, max_frames: int = 16): |
| | super().__init__() |
| | self.dim = dim |
| | self.num_heads = num_heads |
| | self.head_dim = dim // num_heads |
| | self.scale = self.head_dim ** -0.5 |
| | self.max_frames = max_frames |
| | |
| | self.qkv = nn.Linear(dim, dim * 3, bias=False) |
| | self.proj = nn.Linear(dim, dim) |
| | |
| | |
| | self.temporal_pos_embed = nn.Parameter(torch.randn(1, max_frames, dim) * 0.02) |
| | self.frame_type_embed = nn.Parameter(torch.randn(3, dim) * 0.02) |
| | |
| | |
| | self.cross_frame_norm = nn.LayerNorm(dim) |
| | self.cross_frame_mlp = nn.Sequential( |
| | nn.Linear(dim, dim * 2), |
| | nn.GELU(), |
| | nn.Linear(dim * 2, dim) |
| | ) |
| | |
| | def forward(self, current_frames: torch.Tensor, temporal_context: Dict[str, Any], num_current_frames: int = 1): |
| | """ |
| | current_frames: [B*N, H*W, C] where N is number of frames being generated |
| | temporal_context: dict with past frame information |
| | """ |
| | B_times_N, HW, C = current_frames.shape |
| | B = B_times_N // num_current_frames |
| | |
| | if not temporal_context.get("has_context", False): |
| | return current_frames |
| | |
| | |
| | current = current_frames.view(B, num_current_frames, HW, C) |
| | |
| | |
| | past_features = temporal_context["frame_features"] |
| | T, _, H, W = past_features.shape |
| | past_features = past_features.view(T, C, H*W).permute(0, 2, 1) |
| | past_features = past_features.unsqueeze(0).expand(B, -1, -1, -1) |
| | |
| | |
| | all_frames = torch.cat([past_features, current], dim=1) |
| | total_frames = T + num_current_frames |
| | |
| | |
| | pos_ids = torch.arange(total_frames, device=current_frames.device) |
| | pos_embed = self.temporal_pos_embed[:, :total_frames] |
| | |
| | |
| | frame_type_ids = torch.cat([ |
| | torch.zeros(T, device=current_frames.device), |
| | torch.ones(num_current_frames, device=current_frames.device) |
| | ]).long() |
| | type_embed = self.frame_type_embed[frame_type_ids] |
| | |
| | |
| | all_frames = all_frames + pos_embed.unsqueeze(2) + type_embed.unsqueeze(0).unsqueeze(2) |
| | |
| | |
| | all_frames_flat = all_frames.view(B, total_frames * HW, C) |
| | |
| | |
| | qkv = self.qkv(all_frames_flat).reshape(B, -1, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) |
| | q, k, v = qkv[0], qkv[1], qkv[2] |
| | |
| | |
| | attn = (q @ k.transpose(-2, -1)) * self.scale |
| | |
| | |
| | mask = torch.triu(torch.ones(total_frames * HW, total_frames * HW, device=current_frames.device), diagonal=1) |
| | mask = mask.bool() |
| | attn = attn.masked_fill(mask.unsqueeze(0).unsqueeze(0), float('-inf')) |
| | |
| | attn = attn.softmax(dim=-1) |
| | out = (attn @ v).transpose(1, 2).reshape(B, total_frames * HW, C) |
| | |
| | |
| | current_start = T * HW |
| | enhanced_current = out[:, current_start:] |
| | enhanced_current = self.proj(enhanced_current) |
| | |
| | |
| | if num_current_frames > 1: |
| | enhanced_current = enhanced_current.view(B, num_current_frames, HW, C) |
| | for i in range(num_current_frames): |
| | frame_i = enhanced_current[:, i] |
| | |
| | |
| | other_frames = torch.cat([ |
| | enhanced_current[:, :i], |
| | enhanced_current[:, i+1:] |
| | ], dim=1) if num_current_frames > 1 else None |
| | |
| | if other_frames is not None: |
| | cross_context = other_frames.mean(dim=1) |
| | frame_i_norm = self.cross_frame_norm(frame_i + cross_context) |
| | frame_i = frame_i + self.cross_frame_mlp(frame_i_norm) |
| | enhanced_current[:, i] = frame_i |
| | |
| | enhanced_current = enhanced_current.view(B * num_current_frames, HW, C) |
| | |
| | return enhanced_current |
| |
|
| | class FlexibleI2VDiffuser(nn.Module): |
| | """Flexible I2V diffusion model that generates N frames at a time""" |
| | |
| | def __init__( |
| | self, |
| | base_diffusion_model, |
| | feature_dim: int = 512, |
| | temporal_buffer_size: int = 8, |
| | num_attention_heads: int = 8, |
| | max_batch_frames: int = 3 |
| | ): |
| | super().__init__() |
| | self.base_model = base_diffusion_model |
| | self.feature_dim = feature_dim |
| | self.temporal_buffer_size = temporal_buffer_size |
| | self.max_batch_frames = max_batch_frames |
| | |
| | |
| | self.image_encoder = nn.Sequential( |
| | nn.Conv2d(3, feature_dim // 4, 7, padding=3), |
| | nn.GroupNorm(8, feature_dim // 4), |
| | nn.SiLU(), |
| | nn.Conv2d(feature_dim // 4, feature_dim // 2, 3, padding=1, stride=2), |
| | nn.GroupNorm(8, feature_dim // 2), |
| | nn.SiLU(), |
| | nn.Conv2d(feature_dim // 2, feature_dim, 3, padding=1, stride=2), |
| | nn.GroupNorm(8, feature_dim), |
| | nn.SiLU() |
| | ) |
| | |
| | self.latent_encoder = nn.Conv2d( |
| | base_diffusion_model.in_channels, feature_dim, 3, padding=1 |
| | ) |
| | |
| | |
| | self.temporal_attention = FlexibleTemporalAttention( |
| | feature_dim, num_attention_heads, max_batch_frames * 4 |
| | ) |
| | |
| | |
| | self.reference_adapter = nn.Sequential( |
| | nn.Conv2d(feature_dim * 2, feature_dim, 1), |
| | nn.GroupNorm(8, feature_dim), |
| | nn.SiLU() |
| | ) |
| | |
| | self.motion_conditioner = nn.Sequential( |
| | nn.Linear(feature_dim, feature_dim * 2), |
| | nn.GELU(), |
| | nn.Linear(feature_dim * 2, feature_dim) |
| | ) |
| | |
| | |
| | self.frame_consistency_net = nn.Sequential( |
| | nn.Conv3d(feature_dim, feature_dim, (3, 3, 3), padding=(1, 1, 1)), |
| | nn.GroupNorm(8, feature_dim), |
| | nn.SiLU(), |
| | nn.Conv3d(feature_dim, feature_dim, (1, 3, 3), padding=(0, 1, 1)) |
| | ) |
| | |
| | |
| | self.temporal_buffer = TemporalConsistencyBuffer(temporal_buffer_size, feature_dim) |
| | |
| | def encode_reference_image(self, image: torch.Tensor) -> torch.Tensor: |
| | """Encode reference image for I2V conditioning""" |
| | if image.shape[1] == 3: |
| | return self.image_encoder(image) |
| | else: |
| | return self.latent_encoder(image) |
| | |
| | def apply_i2v_conditioning( |
| | self, |
| | current_latents: torch.Tensor, |
| | temporal_context: Dict[str, Any], |
| | num_frames: int = 1 |
| | ) -> torch.Tensor: |
| | """Apply I2V conditioning with flexible frame count""" |
| | |
| | B_times_N, C, H, W = current_latents.shape |
| | B = B_times_N // num_frames |
| | |
| | |
| | current_features = self.latent_encoder(current_latents) |
| | |
| | if not temporal_context.get("has_context", False): |
| | return current_latents |
| | |
| | |
| | current_flat = current_features.flatten(2).transpose(1, 2) |
| | enhanced_features = self.temporal_attention(current_flat, temporal_context, num_frames) |
| | enhanced_features = enhanced_features.transpose(1, 2).reshape(B_times_N, -1, H, W) |
| | |
| | |
| | if temporal_context.get("reference_frame") is not None: |
| | ref_frame = temporal_context["reference_frame"] |
| | ref_features = self.encode_reference_image(ref_frame) |
| | |
| | |
| | ref_features = ref_features.repeat(num_frames, 1, 1, 1) |
| | |
| | |
| | combined_features = torch.cat([enhanced_features, ref_features], dim=1) |
| | conditioned_features = self.reference_adapter(combined_features) |
| | else: |
| | conditioned_features = enhanced_features |
| | |
| | |
| | if num_frames > 1: |
| | |
| | batch_features = conditioned_features.view(B, num_frames, -1, H, W) |
| | batch_features = batch_features.permute(0, 2, 1, 3, 4) |
| | |
| | |
| | consistent_features = self.frame_consistency_net(batch_features) |
| | consistent_features = consistent_features.permute(0, 2, 1, 3, 4) |
| | conditioned_features = consistent_features.reshape(B_times_N, -1, H, W) |
| | |
| | |
| | if temporal_context.get("predicted_motions"): |
| | motions = temporal_context["predicted_motions"][:num_frames] |
| | for i, motion in enumerate(motions): |
| | if motion is not None: |
| | frame_idx = i * B + torch.arange(B, device=current_latents.device) |
| | motion_flat = motion.flatten(2).transpose(1, 2).mean(dim=1) |
| | motion_cond = self.motion_conditioner(motion_flat) |
| | motion_cond = motion_cond.unsqueeze(-1).unsqueeze(-1) |
| | conditioned_features[frame_idx] += motion_cond |
| | |
| | |
| | alpha = 0.4 |
| | enhanced_latents = current_latents + alpha * conditioned_features |
| | |
| | return enhanced_latents |
| | |
| | def forward( |
| | self, |
| | noisy_latents: torch.Tensor, |
| | timestep: torch.Tensor, |
| | text_embeddings: torch.Tensor, |
| | num_frames: int = 1, |
| | use_temporal_consistency: bool = True |
| | ) -> torch.Tensor: |
| | """Forward pass with flexible frame count""" |
| | |
| | if use_temporal_consistency: |
| | |
| | temporal_context = self.temporal_buffer.get_temporal_context() |
| | |
| | |
| | enhanced_latents = self.apply_i2v_conditioning( |
| | noisy_latents, temporal_context, num_frames |
| | ) |
| | else: |
| | enhanced_latents = noisy_latents |
| | |
| | |
| | if text_embeddings.shape[0] != enhanced_latents.shape[0]: |
| | text_embeddings = text_embeddings.repeat(num_frames, 1, 1) |
| | |
| | |
| | noise_pred = self.base_model(enhanced_latents, timestep, text_embeddings) |
| | |
| | return noise_pred |
| | |
| | def update_temporal_buffer(self, latents: torch.Tensor, images: Optional[torch.Tensor] = None, num_frames: int = 1): |
| | """Update temporal buffer with generated frames""" |
| | with torch.no_grad(): |
| | features = self.latent_encoder(latents) |
| | self.temporal_buffer.add_frames(features, latents, images, num_frames) |
| |
|
| | class FlexibleI2VGenerator: |
| | """High-level generator with configurable frame batch sizes""" |
| | |
| | def __init__( |
| | self, |
| | diffusion_model: FlexibleI2VDiffuser, |
| | scheduler, |
| | vae, |
| | device: str = "cuda" |
| | ): |
| | self.model = diffusion_model |
| | self.scheduler = scheduler |
| | self.vae = vae |
| | self.device = device |
| | |
| | |
| | self.image_transform = transforms.Compose([ |
| | transforms.Resize((512, 512)), |
| | transforms.ToTensor(), |
| | transforms.Normalize([0.5], [0.5]) |
| | ]) |
| | |
| | def encode_image(self, image: Union[Image.Image, torch.Tensor]) -> torch.Tensor: |
| | """Encode PIL image or tensor to latent space""" |
| | if isinstance(image, Image.Image): |
| | image = self.image_transform(image).unsqueeze(0).to(self.device) |
| | |
| | with torch.no_grad(): |
| | latent = self.vae.encode(image).latent_dist.sample() |
| | latent = latent * self.vae.config.scaling_factor |
| | |
| | return latent |
| | |
| | def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: |
| | """Decode latents to images""" |
| | with torch.no_grad(): |
| | latents = latents / self.vae.config.scaling_factor |
| | images = self.vae.decode(latents).sample |
| | images = (images + 1.0) / 2.0 |
| | images = torch.clamp(images, 0.0, 1.0) |
| | return images |
| | |
| | @torch.no_grad() |
| | def generate_i2v_sequence( |
| | self, |
| | reference_image: Union[Image.Image, torch.Tensor], |
| | prompt: str, |
| | text_encoder, |
| | tokenizer, |
| | num_frames: int = 16, |
| | frames_per_batch: int = 2, |
| | num_inference_steps: int = 20, |
| | guidance_scale: float = 7.5, |
| | generator: Optional[torch.Generator] = None, |
| | callback=None |
| | ) -> List[torch.Tensor]: |
| | """Generate I2V sequence with configurable batch size""" |
| | |
| | print(f"π¬ Generating {num_frames} frames in batches of {frames_per_batch}") |
| | |
| | |
| | ref_latent = self.encode_image(reference_image) |
| | ref_image_tensor = reference_image if isinstance(reference_image, torch.Tensor) else \ |
| | self.image_transform(reference_image).unsqueeze(0).to(self.device) |
| | |
| | |
| | text_inputs = tokenizer( |
| | prompt, |
| | padding="max_length", |
| | max_length=tokenizer.model_max_length, |
| | truncation=True, |
| | return_tensors="pt" |
| | ) |
| | text_embeddings = text_encoder(text_inputs.input_ids.to(self.device))[0] |
| | |
| | |
| | uncond_tokens = [""] |
| | uncond_inputs = tokenizer( |
| | uncond_tokens, |
| | padding="max_length", |
| | max_length=tokenizer.model_max_length, |
| | return_tensors="pt" |
| | ) |
| | uncond_embeddings = text_encoder(uncond_inputs.input_ids.to(self.device))[0] |
| | |
| | |
| | self.model.temporal_buffer = TemporalConsistencyBuffer( |
| | self.model.temporal_buffer_size, |
| | self.model.feature_dim |
| | ) |
| | self.model.update_temporal_buffer(ref_latent, ref_image_tensor, 1) |
| | |
| | generated_frames = [ref_latent] |
| | latent_shape = ref_latent.shape |
| | |
| | |
| | frames_generated = 1 |
| | |
| | while frames_generated < num_frames: |
| | |
| | remaining_frames = num_frames - frames_generated |
| | current_batch_size = min(frames_per_batch, remaining_frames) |
| | |
| | print(f"π― Generating frames {frames_generated+1}-{frames_generated+current_batch_size}") |
| | |
| | |
| | batch_latents = torch.randn( |
| | (current_batch_size, *latent_shape[1:]), |
| | generator=generator, |
| | device=self.device, |
| | dtype=text_embeddings.dtype |
| | ) |
| | |
| | |
| | batch_text_embeddings = torch.cat([ |
| | uncond_embeddings.repeat(current_batch_size, 1, 1), |
| | text_embeddings.repeat(current_batch_size, 1, 1) |
| | ]) |
| | |
| | |
| | self.scheduler.set_timesteps(num_inference_steps, device=self.device) |
| | timesteps = self.scheduler.timesteps |
| | |
| | |
| | for i, t in enumerate(timesteps): |
| | |
| | latent_model_input = torch.cat([batch_latents] * 2) |
| | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) |
| | |
| | |
| | noise_pred = self.model( |
| | latent_model_input, |
| | t, |
| | batch_text_embeddings, |
| | num_frames=current_batch_size, |
| | use_temporal_consistency=True |
| | ) |
| | |
| | |
| | noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) |
| | noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) |
| | |
| | |
| | batch_latents = self.scheduler.step(noise_pred, t, batch_latents).prev_sample |
| | |
| | if callback: |
| | callback(i, t, batch_latents) |
| | |
| | |
| | batch_images = self.decode_latents(batch_latents) |
| | self.model.update_temporal_buffer(batch_latents, batch_images, current_batch_size) |
| | |
| | |
| | for j in range(current_batch_size): |
| | generated_frames.append(batch_latents[j:j+1]) |
| | |
| | frames_generated += current_batch_size |
| | print(f"β
Generated {current_batch_size} frames") |
| | |
| | return generated_frames |
| | |
| | def generate_with_stepping_strategy( |
| | self, |
| | reference_image: Union[Image.Image, torch.Tensor], |
| | prompt: str, |
| | text_encoder, |
| | tokenizer, |
| | total_frames: int = 24, |
| | stepping_pattern: List[int] = [1, 2, 3, 2, 1], |
| | **kwargs |
| | ) -> List[torch.Tensor]: |
| | """Generate with dynamic stepping pattern""" |
| | |
| | all_frames = [] |
| | frames_generated = 0 |
| | step_idx = 0 |
| | |
| | while frames_generated < total_frames: |
| | |
| | current_step = stepping_pattern[step_idx % len(stepping_pattern)] |
| | remaining = total_frames - frames_generated |
| | actual_step = min(current_step, remaining) |
| | |
| | print(f"π Step {step_idx + 1}: Generating {actual_step} frames") |
| | |
| | |
| | if frames_generated == 0: |
| | |
| | frames = self.generate_i2v_sequence( |
| | reference_image=reference_image, |
| | prompt=prompt, |
| | text_encoder=text_encoder, |
| | tokenizer=tokenizer, |
| | num_frames=actual_step + 1, |
| | frames_per_batch=actual_step, |
| | **kwargs |
| | ) |
| | all_frames.extend(frames) |
| | frames_generated += len(frames) |
| | else: |
| | |
| | last_frame_latent = all_frames[-1] |
| | last_frame_image = self.decode_latents(last_frame_latent) |
| | |
| | frames = self.generate_i2v_sequence( |
| | reference_image=last_frame_image, |
| | prompt=prompt, |
| | text_encoder=text_encoder, |
| | tokenizer=tokenizer, |
| | num_frames=actual_step + 1, |
| | frames_per_batch=actual_step, |
| | **kwargs |
| | ) |
| | all_frames.extend(frames[1:]) |
| | frames_generated += len(frames) - 1 |
| | |
| | step_idx += 1 |
| | |
| | return all_frames[:total_frames] |
| |
|
| | |
| | class TemporalMiddleware: |
| | """Middleware layer for external AI control""" |
| | |
| | def __init__(self): |
| | self.prompt_scheduler = None |
| | self.controlnet_adapter = None |
| | self.audio_sync = None |
| | |
| | def intercept_temporal_state(self, temporal_context: Dict, frame_idx: int) -> Dict: |
| | """Hook for external manipulation of temporal state""" |
| | |
| | |
| | if self.prompt_scheduler: |
| | new_prompt = self.prompt_scheduler.get_prompt_at_frame(frame_idx) |
| | temporal_context["dynamic_prompt"] = new_prompt |
| | |
| | |
| | if self.controlnet_adapter: |
| | control_inputs = self.controlnet_adapter.get_control_at_frame(frame_idx) |
| | temporal_context["control_inputs"] = control_inputs |
| | |
| | |
| | if self.audio_sync: |
| | audio_features = self.audio_sync.get_features_at_frame(frame_idx) |
| | temporal_context["audio_conditioning"] = audio_features |
| | |
| | return temporal_context |
| |
|
| | def example_usage(): |
| | """Example of flexible I2V generation with middleware""" |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | |
| | print("π Flexible Batch I2V Generator with Temporal Consistency - IMPLEMENTED!") |
| | print("π Ready for infinite frame generation with external AI control!") |
| |
|
| | if __name__ == "__main__": |
| | example_usage() |
| |
|
| |
|
| |
|