From 7ea7e51883cfac6f0a72d8ec6ba9218bef66e0b4 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Tue, 11 Mar 2025 15:38:49 -0400 Subject: [PATCH] sans profiling --- stream_sampler.py | 147 ++-------------------------------------------- 1 file changed, 6 insertions(+), 141 deletions(-) diff --git a/stream_sampler.py b/stream_sampler.py index f61328c..e9cb1a2 100644 --- a/stream_sampler.py +++ b/stream_sampler.py @@ -4,134 +4,6 @@ import comfy.samplers import random import time import os -import json -from datetime import datetime -from functools import wraps -import cProfile -import pstats -import io - -# Simple profiling setup - SINGLE FILE -PROFILE_FILE = "stream_sampler_profile.json" -PROFILE_DATA = [] - -def profile_time(func): - @wraps(func) - def wrapper(*args, **kwargs): - # Start timing - start_time = time.time() - - # Check memory before - if torch.cuda.is_available(): - torch.cuda.synchronize() - mem_before = torch.cuda.memory_allocated() / (1024 * 1024) - - # Execute the function - result = func(*args, **kwargs) - - # End timing - if torch.cuda.is_available(): - torch.cuda.synchronize() - end_time = time.time() - - # Check memory after - if torch.cuda.is_available(): - mem_after = torch.cuda.memory_allocated() / (1024 * 1024) - mem_diff = mem_after - mem_before - else: - mem_diff = 0 - - exec_time = end_time - start_time - - # Get shape info if available - shape_info = "unknown" - if len(args) > 1 and hasattr(args[1], 'shape'): # If this is sample(), args[1] is noise - shape_info = str(args[1].shape) - - # Log results - entry = { - 'function': func.__name__, - 'execution_time': exec_time, - 'memory_diff_mb': mem_diff, - 'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), - 'shape': shape_info - } - - PROFILE_DATA.append(entry) - - # Print to console - print(f"PROFILE: {func.__name__} - Time: {exec_time:.4f}s, Memory: {mem_diff:.2f}MB") - - # Save to file after each run - save_profile_data() - - return result - return wrapper - - -def profile_cprofile(func): - @wraps(func) - def wrapper(*args, **kwargs): - # Create profiler - pr = cProfile.Profile() - pr.enable() - - # Execute the function - result = func(*args, **kwargs) - - # Disable profiler - pr.disable() - - # Get stats - s = io.StringIO() - ps = pstats.Stats(pr, stream=s).sort_stats('cumulative') - ps.print_stats(20) - - # Add to the most recent profile entry - if PROFILE_DATA: - PROFILE_DATA[-1]['cprofile_data'] = s.getvalue() - - return result - return wrapper - - -def save_profile_data(): - """Save all profiling data to a single JSON file""" - if not PROFILE_DATA: - return - - # Convert to serializable format - json_data = [] - for entry in PROFILE_DATA: - serializable_entry = { - 'function': entry['function'], - 'execution_time': float(entry['execution_time']), - 'memory_diff_mb': float(entry['memory_diff_mb']), - 'timestamp': entry['timestamp'], - 'shape': entry.get('shape', 'unknown') - } - - if 'cprofile_data' in entry: - # Only store recent cprofile data to keep file size manageable - if len(json_data) < 10 or len(json_data) % 10 == 0: - serializable_entry['cprofile_data'] = entry['cprofile_data'] - - json_data.append(serializable_entry) - - # Write to a single file - with open(PROFILE_FILE, 'w') as f: - json.dump(json_data, f, indent=2) - - # Calculate stats - sample_times = [x['execution_time'] for x in PROFILE_DATA if x['function'] == 'sample'] - if sample_times: - avg_time = sum(sample_times) / len(sample_times) - min_time = min(sample_times) - max_time = max(sample_times) - print(f"PROFILE SUMMARY: {len(sample_times)} runs, Avg: {avg_time:.4f}s, Min: {min_time:.4f}s, Max: {max_time:.4f}s") - - print(f"Profiling data saved to {PROFILE_FILE}") - class StreamBatchSampler: @@ -169,8 +41,6 @@ class StreamBatchSampler: self.working_buffer = None self.output_buffer = None - @profile_time - @profile_cprofile def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None): """Sample with staggered batch denoising steps - Optimized version""" extra_args = {} if extra_args is None else extra_args @@ -279,16 +149,12 @@ class StreamBatchSampler: return x_0_pred_out - @profile_time def update(self, num_steps=4): """Create sampler with specified settings""" self.num_steps = num_steps sampler = comfy.samplers.KSAMPLER(self.sample) return (sampler,) -# Print setup info when module is imported -print(f"StreamBatchSampler profiling enabled. Results will be saved to {PROFILE_FILE}") - class StreamScheduler: @@ -323,7 +189,7 @@ class StreamScheduler: try: t_index_list = [int(t.strip()) for t in t_index_list.split(",")] except ValueError as e: - print(f"Error parsing timesteps: {e}. Using default [32,45]") + t_index_list = [32, 45] # Create full schedule using normal scheduler @@ -334,7 +200,7 @@ class StreamScheduler: selected_sigmas = [] for t in sorted(t_index_list, reverse=True): # Sort in reverse to go from high noise to low if t < 0 or t >= num_inference_steps: - print(f"Warning: timestep {t} out of range [0,{num_inference_steps}), skipping") + continue selected_sigmas.append(float(full_sigmas[t])) @@ -387,7 +253,7 @@ class StreamFrameBuffer: if is_empty_latent: self.is_txt2img_mode = True - print(f"[StreamFrameBuffer] Detected empty latent for text-to-image mode") + # Create empty latents with correct dimensions for txt2img # Get dimensions from latent dict height = latent.get("height", 512) @@ -400,7 +266,7 @@ class StreamFrameBuffer: # Create zero tensor with correct shape x = torch.zeros((4, latent_height, latent_width), device=comfy.model_management.get_torch_device()) - print(f"[StreamFrameBuffer] Created empty latent with shape: {x.shape}") + elif x.dim() == 4: # [B,C,H,W] self.is_txt2img_mode = False x = x.squeeze(0) # Remove batch dimension -> [C,H,W] @@ -417,18 +283,17 @@ class StreamFrameBuffer: if self.is_txt2img_mode or not self.is_initialized: # Optimization: Use broadcasting to fill buffer with copies self.frame_buffer[:] = x.unsqueeze(0) # Broadcast x to [buffer_size, C, H, W] - print(f"[StreamFrameBuffer] Initialized buffer with {self.buffer_size} copies of frame") + self.is_initialized = True self.buffer_pos = 0 else: # Add new frame to buffer using ring buffer logic self.frame_buffer[self.buffer_pos] = x # In-place update - print(f"[StreamFrameBuffer] Added new frame to buffer at position {self.buffer_pos}") self.buffer_pos = (self.buffer_pos + 1) % self.buffer_size # Circular increment # Optimization: frame_buffer is already a tensor batch, no need to stack batch = self.frame_buffer - print(f"[StreamFrameBuffer] Created batch with shape: {batch.shape}") + # Return as latent dict with preserved dimensions result = {"samples": batch}