Create poc.py

poc.py

@@ -0,0 +1,656 @@
+
+"""
+Flexible Batch I2V Generator with Temporal Consistency
+Generates N frames at a time (1, 2, 3, etc.) while maintaining temporal consistency
+Optimized for Image-to-Video models with reference frame initialization
+"""
+
+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)  # Store actual frames for I2V
+        self.motion_vectors = deque(maxlen=buffer_size-1)
+        self.temporal_weights = deque(maxlen=buffer_size)  # Importance weights
+        
+    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
+            
+            # Calculate motion vector if we have previous frames
+            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 newer frames more heavily
+            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}
+        
+        # Get most recent frames up to num_context_frames
+        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:]
+        
+        # Stack with attention to batch dimension
+        stacked_features = torch.cat(recent_features, dim=0)  # [T, C, H, W]
+        stacked_latents = torch.cat(recent_latents, dim=0)
+        weights = torch.tensor(recent_weights, device=stacked_features.device)
+        
+        # Predict motion for next frames
+        predicted_motions = []
+        if len(self.motion_vectors) >= 2:
+            # Multi-step motion prediction
+            recent_motions = list(self.motion_vectors)[-3:]  # Last 3 motions
+            for step in range(1, 4):  # Predict up to 3 steps ahead
+                if len(recent_motions) >= 2:
+                    # Weighted motion extrapolation
+                    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)
+        
+        # Learnable temporal positional embeddings
+        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)  # past, current, future
+        
+        # Cross-frame interaction
+        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
+        
+        # Reshape current frames
+        current = current_frames.view(B, num_current_frames, HW, C)
+        
+        # Get temporal context
+        past_features = temporal_context["frame_features"]  # [T, C, H, W]
+        T, _, H, W = past_features.shape
+        past_features = past_features.view(T, C, H*W).permute(0, 2, 1)  # [T, H*W, C]
+        past_features = past_features.unsqueeze(0).expand(B, -1, -1, -1)  # [B, T, H*W, C]
+        
+        # Combine all frames (past + current)
+        all_frames = torch.cat([past_features, current], dim=1)  # [B, T+N, H*W, C]
+        total_frames = T + num_current_frames
+        
+        # Add positional embeddings
+        pos_ids = torch.arange(total_frames, device=current_frames.device)
+        pos_embed = self.temporal_pos_embed[:, :total_frames]  # [1, T+N, C]
+        
+        # Add frame type embeddings (past=0, current=1, future=2)
+        frame_type_ids = torch.cat([
+            torch.zeros(T, device=current_frames.device),  # past frames
+            torch.ones(num_current_frames, device=current_frames.device)  # current frames
+        ]).long()
+        type_embed = self.frame_type_embed[frame_type_ids]  # [T+N, C]
+        
+        # Apply embeddings
+        all_frames = all_frames + pos_embed.unsqueeze(2) + type_embed.unsqueeze(0).unsqueeze(2)
+        
+        # Flatten for attention
+        all_frames_flat = all_frames.view(B, total_frames * HW, C)
+        
+        # Multi-head attention
+        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]
+        
+        # Temporal attention with causality mask for current frames
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        
+        # Create causal mask - current frames can see past + themselves, but not future
+        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)
+        
+        # Extract only current frame features
+        current_start = T * HW
+        enhanced_current = out[:, current_start:]  # [B, N*H*W, C]
+        enhanced_current = self.proj(enhanced_current)
+        
+        # Cross-frame interaction within current batch
+        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]  # [B, H*W, C]
+                
+                # Interact with other frames in current batch
+                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)  # [B, H*W, C]
+                    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
+        
+        # Enhanced feature extraction for I2V
+        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
+        )
+        
+        # Flexible temporal attention
+        self.temporal_attention = FlexibleTemporalAttention(
+            feature_dim, num_attention_heads, max_batch_frames * 4
+        )
+        
+        # I2V specific components
+        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)
+        )
+        
+        # Multi-frame consistency
+        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))
+        )
+        
+
+        # Initialize temporal buffer
+        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:  # RGB image
+            return self.image_encoder(image)
+        else:  # Already encoded latent
+            return self.latent_encoder(image)
+    
+    def apply_i2v_conditioning(
+        self,
+        current_latents: torch.Tensor,  # [B*N, C, H, W]
+        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
+        
+        # Extract features from current latents
+        current_features = self.latent_encoder(current_latents)  # [B*N, F, H, W]
+        
+        if not temporal_context.get("has_context", False):
+            return current_latents
+        
+        # Apply temporal attention
+        current_flat = current_features.flatten(2).transpose(1, 2)  # [B*N, H*W, F]
+        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)
+        
+        # Reference frame conditioning for I2V
+        if temporal_context.get("reference_frame") is not None:
+            ref_frame = temporal_context["reference_frame"]
+            ref_features = self.encode_reference_image(ref_frame)
+            
+            # Broadcast reference to all current frames
+            ref_features = ref_features.repeat(num_frames, 1, 1, 1)
+            
+            # Combine with current features
+            combined_features = torch.cat([enhanced_features, ref_features], dim=1)
+            conditioned_features = self.reference_adapter(combined_features)
+        else:
+            conditioned_features = enhanced_features
+        
+        # Multi-frame consistency for batch generation
+        if num_frames > 1:
+            # Reshape for 3D convolution
+            batch_features = conditioned_features.view(B, num_frames, -1, H, W)
+            batch_features = batch_features.permute(0, 2, 1, 3, 4)  # [B, C, T, H, W]
+            
+            # Apply 3D consistency
+            consistent_features = self.frame_consistency_net(batch_features)
+            consistent_features = consistent_features.permute(0, 2, 1, 3, 4)  # [B, T, C, H, W]
+            conditioned_features = consistent_features.reshape(B_times_N, -1, H, W)
+        
+        # Motion conditioning
+        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)  # [B, F]
+                    motion_cond = self.motion_conditioner(motion_flat)  # [B, F]
+                    motion_cond = motion_cond.unsqueeze(-1).unsqueeze(-1)  # [B, F, 1, 1]
+                    conditioned_features[frame_idx] += motion_cond
+        
+        # Blend with original latents
+        alpha = 0.4  # I2V conditioning strength
+        enhanced_latents = current_latents + alpha * conditioned_features
+        
+        return enhanced_latents
+    
+    def forward(
+        self,
+        noisy_latents: torch.Tensor,  # [B*N, C, H, W]
+        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:
+            # Get temporal context
+            temporal_context = self.temporal_buffer.get_temporal_context()
+            
+            # Apply I2V conditioning
+            enhanced_latents = self.apply_i2v_conditioning(
+                noisy_latents, temporal_context, num_frames
+            )
+        else:
+            enhanced_latents = noisy_latents
+        
+        # Expand text embeddings for multiple frames
+        if text_embeddings.shape[0] != enhanced_latents.shape[0]:
+            text_embeddings = text_embeddings.repeat(num_frames, 1, 1)
+        
+        # Run base diffusion model
+        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,  # For encoding/decoding images
+        device: str = "cuda"
+    ):
+        self.model = diffusion_model
+        self.scheduler = scheduler  
+        self.vae = vae
+        self.device = device
+        
+        # Image preprocessing
+        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,  # This is the key parameter!
+        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}")
+        
+        # Encode reference image
+        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)
+        
+        # Encode text prompt
+        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]
+        
+        # Prepare unconditional embeddings
+        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]
+        
+        # Reset temporal buffer and add reference frame
+        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
+        
+        # Generate in flexible batches
+        frames_generated = 1  # Start with reference frame
+        
+        while frames_generated < num_frames:
+            # Calculate current batch size
+            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}")
+            
+            # Initialize noise for current batch
+            batch_latents = torch.randn(
+                (current_batch_size, *latent_shape[1:]),
+                generator=generator,
+                device=self.device,
+                dtype=text_embeddings.dtype
+            )
+            
+            # Prepare embeddings for batch
+            batch_text_embeddings = torch.cat([
+                uncond_embeddings.repeat(current_batch_size, 1, 1),
+                text_embeddings.repeat(current_batch_size, 1, 1)
+            ])
+            
+            # Set scheduler timesteps
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            timesteps = self.scheduler.timesteps
+            
+            # Denoising loop for current batch
+            for i, t in enumerate(timesteps):
+                # Expand for classifier-free guidance
+                latent_model_input = torch.cat([batch_latents] * 2)
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                
+                # Predict noise with temporal consistency
+                noise_pred = self.model(
+                    latent_model_input,
+                    t,
+                    batch_text_embeddings,
+                    num_frames=current_batch_size,
+                    use_temporal_consistency=True
+                )
+                
+                # Classifier-free guidance
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                
+                # Scheduler step
+                batch_latents = self.scheduler.step(noise_pred, t, batch_latents).prev_sample
+                
+                if callback:
+                    callback(i, t, batch_latents)
+            
+            # Update temporal buffer with generated batch
+            batch_images = self.decode_latents(batch_latents)
+            self.model.update_temporal_buffer(batch_latents, batch_images, current_batch_size)
+            
+            # Add to results
+            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],  # Variable batch sizes
+        **kwargs
+    ) -> List[torch.Tensor]:
+        """Generate with dynamic stepping pattern"""
+        
+        all_frames = []
+        frames_generated = 0
+        step_idx = 0
+        
+        while frames_generated < total_frames:
+            # Get current step size
+            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")
+            
+            # Generate batch
+            if frames_generated == 0:
+                # First generation includes reference
+                frames = self.generate_i2v_sequence(
+                    reference_image=reference_image,
+                    prompt=prompt,
+                    text_encoder=text_encoder,
+                    tokenizer=tokenizer,
+                    num_frames=actual_step + 1,  # +1 for reference
+                    frames_per_batch=actual_step,
+                    **kwargs
+                )
+                all_frames.extend(frames)
+                frames_generated += len(frames)
+            else:
+                # Continue from last frame
+                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:])  # Skip reference (duplicate)
+                frames_generated += len(frames) - 1
+            
+            step_idx += 1
+        
+        return all_frames[:total_frames]  # Ensure exact frame count
+
+# Example usage
+def example_usage():
+    """Example of flexible I2V generation"""
+    
+    # Load your models (example)
+    #
+    base_model, scheduler, vae, text_encoder, tokenizer = load_models()
+    
+    # Create flexible I2V model
+    i2v_model = FlexibleI2VDiffuser(
+        base_diffusion_model=i2v_model,
+        feature_dim=512,
+        temporal_buffer_size=8,
+        max_batch_frames=3
+    )
+    
+    # Create generator
+    generator = FlexibleI2VGenerator(
+        diffusion_model=i2v_model,
+        scheduler=scheduler,
+        vae=vae,
+        device="cuda"
+    )
+    
+    # Load reference image
+    reference_image = Image.open("reference.jpg")
+    
+    # Strategy 1: Fixed batch size
+    frames_fixed = generator.generate_i2v_sequence(
+        reference_image=reference_image,
+        prompt="A cat walking in a garden",
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        num_frames=16,
+        frames_per_batch=2,  # Generate 2 frames at a time
+        num_inference_steps=20
+    )
+    
+    # Strategy 2: Variable stepping pattern
+    frames_variable = generator.generate_with_stepping_strategy(
+        reference_image=reference_image,
+        prompt="A cat walking in a garden", 
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        total_frames=24,
+        stepping_pattern=[1, 2, 3, 2, 1],  # Start slow, ramp up, slow down
+        num_inference_steps=20
+    )
+    
+    print(f"🎉 Generated {len(frames_fixed)} frames with fixed batching")
+    print(f"🎉 Generated {len(frames_variable)} frames with variable stepping")
+
+if __name__ == "__main__":
+    example_usage()