peterlllmm/infinitetalk

1"""2GPU Memory Manager for InfiniteTalk3Handles memory monitoring, cleanup, and optimization4"""5 6import torch7import logging8from typing import Optional9 10logging.basicConfig(level=logging.INFO)11logger = logging.getLogger(__name__)12 13 14class GPUManager:15    """Manages GPU memory usage and optimization"""16 17    def __init__(self, max_memory_gb=65):18        """19        Initialize GPU Manager20 21        Args:22            max_memory_gb: Maximum memory threshold in GB (default 65GB for 70GB H200)23        """24        self.max_memory_bytes = max_memory_gb * 1024 ** 325        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")26 27    def get_memory_usage(self):28        """29        Get current GPU memory usage30 31        Returns:32            dict with allocated, reserved, and free memory in GB33        """34        if not torch.cuda.is_available():35            return {"allocated": 0, "reserved": 0, "free": 0}36 37        allocated = torch.cuda.memory_allocated() / 1024 ** 338        reserved = torch.cuda.memory_reserved() / 1024 ** 339        total = torch.cuda.get_device_properties(0).total_memory / 1024 ** 340        free = total - allocated41 42        return {43            "allocated": round(allocated, 2),44            "reserved": round(reserved, 2),45            "free": round(free, 2),46            "total": round(total, 2)47        }48 49    def print_memory_usage(self, prefix=""):50        """Print current memory usage"""51        usage = self.get_memory_usage()52        logger.info(53            f"{prefix}GPU Memory - "54            f"Allocated: {usage['allocated']}GB, "55            f"Reserved: {usage['reserved']}GB, "56            f"Free: {usage['free']}GB"57        )58 59    def check_memory_threshold(self):60        """61        Check if memory usage exceeds threshold62 63        Returns:64            bool: True if within safe limits, False if exceeded65        """66        if not torch.cuda.is_available():67            return True68 69        allocated = torch.cuda.memory_allocated()70 71        if allocated > self.max_memory_bytes:72            logger.warning(73                f"Memory threshold exceeded! "74                f"Allocated: {allocated / 1024**3:.2f}GB, "75                f"Threshold: {self.max_memory_bytes / 1024**3:.2f}GB"76            )77            return False78 79        return True80 81    def cleanup(self):82        """Perform garbage collection and CUDA cache cleanup"""83        import gc84 85        gc.collect()86        if torch.cuda.is_available():87            torch.cuda.empty_cache()88            torch.cuda.synchronize()89 90        logger.info("GPU memory cleaned up")91        self.print_memory_usage("After cleanup - ")92 93    def optimize_model_for_inference(self, model):94        """95        Apply optimizations to model for inference96 97        Args:98            model: PyTorch model to optimize99 100        Returns:101            Optimized model102        """103        model.eval()104 105        # Enable gradient checkpointing if available106        if hasattr(model, "enable_gradient_checkpointing"):107            model.enable_gradient_checkpointing()108 109        # Use FP16 for inference to save memory110        if torch.cuda.is_available() and hasattr(model, "half"):111            logger.info("Converting model to FP16")112            model = model.half()113 114        return model115 116    def enable_memory_efficient_attention(self):117        """Enable memory-efficient attention mechanisms"""118        try:119            import xformers120 121            logger.info("xformers available - memory efficient attention enabled")122            return True123        except ImportError:124            logger.warning("xformers not available - using standard attention")125            return False126 127    def estimate_inference_memory(self, resolution="480p", duration_seconds=10):128        """129        Estimate memory requirements for inference130 131        Args:132            resolution: Video resolution (480p or 720p)133            duration_seconds: Video duration in seconds134 135        Returns:136            Estimated memory in GB137        """138        base_memory = 20  # Base model memory139 140        if resolution == "720p":141            per_second_memory = 1.5142        else:  # 480p143            per_second_memory = 0.8144 145        estimated = base_memory + (duration_seconds * per_second_memory)146 147        logger.info(148            f"Estimated memory for {resolution} video ({duration_seconds}s): "149            f"{estimated:.2f}GB"150        )151 152        return estimated153 154    def should_use_chunking(self, video_duration, resolution="480p"):155        """156        Determine if chunked processing should be used157 158        Args:159            video_duration: Duration in seconds160            resolution: Video resolution161 162        Returns:163            bool: True if chunking recommended164        """165        estimated_memory = self.estimate_inference_memory(resolution, video_duration)166 167        # Use chunking if estimated memory exceeds 50GB168        return estimated_memory > 50169 170    def get_optimal_chunk_size(self, resolution="480p"):171        """172        Get optimal chunk size for video processing173 174        Args:175            resolution: Video resolution176 177        Returns:178            Optimal chunk size in seconds179        """180        if resolution == "720p":181            return 10  # 10 second chunks for 720p182        else:183            return 15  # 15 second chunks for 480p184 185    @staticmethod186    def calculate_duration_for_zerogpu(video_duration, resolution="480p"):187        """188        Calculate ZeroGPU duration parameter189 190        Args:191            video_duration: Duration of video in seconds192            resolution: Video resolution193 194        Returns:195            Recommended duration for @spaces.GPU decorator196        """197        base_time = 60  # Base time for model loading198 199        # Processing time per second of video200        if resolution == "720p":201            processing_rate = 3.5202        else:  # 480p203            processing_rate = 2.5204 205        # Add safety margin of 1.2x206        estimated_time = base_time + (video_duration * processing_rate)207        duration = int(estimated_time * 1.2)208 209        # Cap at 300 seconds for free tier (300s ZeroGPU = 10 min real time)210        duration = min(duration, 300)211 212        logger.info(213            f"Calculated ZeroGPU duration: {duration}s for "214            f"{video_duration}s {resolution} video"215        )216 217        return duration218 219 220# Global instance221gpu_manager = GPUManager()222