# ComfyUI_DWposeDeluxe/scripts/memplot.py import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import os import json import math from . import logger def create_memplot(log_path): if not os.path.exists(log_path): return steps = [] ram_data = [] vram_data = [] sections = [] metadata = {} current_label = "Init" ram_total = 0.0 vram_total = 0.0 sections.append({'start': 0, 'label': current_label}) try: with open(log_path, "r") as f: lines = f.readlines() for line in lines: line = line.strip() if not line: continue if line.startswith("# Metadata:"): try: json_str = line.split(":", 1)[1].strip() metadata = json.loads(json_str) except: pass elif line.startswith("# [Loop:"): try: content = line[8:-1] parts = [p.strip() for p in content.split(",")] new_label = parts[0] for p in parts: if "RAM_Total" in p: val = float(p.split(":")[1]) if val > 0: ram_total = val if "VRAM_Total" in p: val = float(p.split(":")[1]) if val > 0: vram_total = val if len(steps) > sections[-1]['start']: current_label = new_label sections.append({'start': len(steps), 'label': current_label}) else: sections[-1]['label'] = new_label current_label = new_label except Exception as e: print(f"Error parsing loop header: {e}") else: try: parts = line.split(",") steps.append(len(steps) + 1) ram_data.append(float(parts[1])) vram_data.append(float(parts[2])) except: pass except Exception as e: print(f"Failed to parse log file: {e}") return if not steps: return # -- Plotting -- plt.style.use('dark_background') fig, ax1 = plt.subplots(figsize=(12, 6)) title_str = "Memory Log" if metadata: provider = metadata.get('provider', '?') batch = metadata.get('batch_size', '?') w = metadata.get('width', '?') h = metadata.get('height', '?') poses = metadata.get('poses_to_detect', '?') title_str = f"Memory Log | {provider} | Batch: {batch} | {w}x{h} | Poses: {poses}" plt.title(title_str, color='white', fontsize=10, pad=20) # 1. Calculate Deltas start_ram = ram_data[0] delta_ram = [max(0, x - start_ram) for x in ram_data] # Colors color_ram = '#FF0000' # Red color_vram = '#0064C8' # Cyan/Blue # 2. Plot RAM (Left Axis) ax1.set_xlabel('Steps') ax1.set_ylabel('RAM (GB)', color='white') line1, = ax1.plot(steps, ram_data, color=color_ram, linewidth=2.0, label='System RAM') ax1.tick_params(axis='y', labelcolor=color_ram) # Scale RAM Y-Axis: Next full integer above total, plus 10% buffer max_ram_val = ram_total if ram_total > 0 else (max(ram_data) if ram_data else 10) ylim_ram_calculated = math.ceil(max_ram_val) if ylim_ram_calculated == max_ram_val: ylim_ram_calculated += 1 ylim_ram_buffered = ylim_ram_calculated * 1.1 # Add 10% buffer for headroom ax1.set_ylim(0, ylim_ram_buffered) # Plot RAM Limit Line if ram_total > 0: ax1.axhline(y=ram_total, color=color_ram, linestyle='--', alpha=0.5, linewidth=1.0) ax1.text(steps[0], ram_total, f' Total: {ram_total:.1f}GB', color=color_ram, va='bottom', fontsize=8) ax1_delta = ax1.twinx() ax1_delta.set_axis_off() max_dram = max(delta_ram) if delta_ram else 0 if max_dram > 0: ax1_delta.set_ylim(0, max_dram * 1.1) line1_d, = ax1_delta.plot(steps, delta_ram, color=color_ram, linewidth=1.5, linestyle=':', label='Process RAM Delta') # 4. Plot VRAM (Right Axis) ax2 = ax1.twinx() ax2.set_ylabel('VRAM (GB)', color='white') ax2.tick_params(axis='y', labelcolor=color_vram) # Determine Y-Limit for VRAM max_vram_val = vram_total if vram_total > 0 else max(vram_data) if vram_data else 0 ylim_vram_calculated = math.ceil(max_vram_val) if ylim_vram_calculated == max_vram_val: ylim_vram_calculated += 1 if ylim_vram_calculated == 0: ylim_vram_calculated = 1 # Fallback if everything is 0 ylim_vram_buffered = ylim_vram_calculated * 1.1 # Add 10% buffer for headroom ax2.set_ylim(0, ylim_vram_buffered) # Handling VRAM logic (Issue 1 & 2) lines_for_legend = [line1, line1_d] names_for_legend = ['System RAM', 'Process RAM Delta'] if vram_total == 0 and (not vram_data or max(vram_data) == 0): # Issue 1: Force only 0 tick if no VRAM ax2.set_ylim(0, 1) # Arbitrary low range ax2.set_yticks([0]) # Issue 2: Don't plot lines else: # Normal VRAM plotting ax2.set_ylim(0, ylim_vram_buffered) start_vram = vram_data[0] if vram_data else 0 delta_vram = [max(0, x - start_vram) for x in vram_data] line2, = ax2.plot(steps, vram_data, color=color_vram, linewidth=2.0, label='System VRAM') if vram_total > 0: ax2.axhline(y=vram_total, color=color_vram, linestyle='--', alpha=0.5, linewidth=1.0) ax2.text(steps[-1] if steps else 0, vram_total, f'Total: {vram_total:.1f}GB ', color=color_vram, va='bottom', ha='right', fontsize=8) ax2_delta = ax1.twinx() ax2_delta.set_axis_off() max_dvram = max(delta_vram) if delta_vram else 0 if max_dvram > 0: ax2_delta.set_ylim(0, max_dvram * 1.1) line2_d, = ax2_delta.plot(steps, delta_vram, color=color_vram, linewidth=1.5, linestyle=':', label='Process VRAM Delta') lines_for_legend.extend([line2, line2_d]) names_for_legend.extend(['System VRAM', 'Process VRAM Delta']) # Legend ax1.legend(lines_for_legend, names_for_legend, loc='lower right', facecolor='black', framealpha=0.7) # 4. Custom X-Axis Ticks (Issue 3) custom_ticks = [] custom_labels = [] # Calculate text position for labels y_text_pos = ax1.get_ylim()[1] * 0.95 # Start Divider ax1.axvline(x=0, color='white', linestyle='--', alpha=0.3) # End Divider if steps: ax1.axvline(x=steps[-1], color='white', linestyle='--', alpha=0.3) for i, sec in enumerate(sections): start_idx = sec['start'] end_idx = sections[i+1]['start'] if i+1 < len(sections) else len(steps) if end_idx > start_idx: duration = end_idx - start_idx # Divider Line if start_idx > 0: ax1.axvline(x=start_idx, color='white', linestyle='--', alpha=0.3) # Section Label mid_point = (start_idx + end_idx) / 2 ax1.text(mid_point, y_text_pos, sec['label'], color='white', ha='center', va='top', fontsize=9, fontweight='bold', bbox=dict(facecolor='black', alpha=0.6, edgecolor='none', pad=2)) # Tick Generation # Always include 0 (local) and Duration (local) # Intermediate ticks: max 10 total step_size = max(1, duration // 8) local_tick_vals = list(range(0, duration + 1, step_size)) if local_tick_vals[-1] != duration: local_tick_vals.append(duration) # Limit to ~10 just in case if len(local_tick_vals) > 12: step_size = max(1, duration // 5) local_tick_vals = list(range(0, duration + 1, step_size)) if local_tick_vals[-1] != duration: local_tick_vals.append(duration) for lt in local_tick_vals: global_pos = start_idx + lt # Offset logic for boundaries if lt == 0: # Offset first tick slightly right to avoid hitting the divider global_pos += (duration * 0.03) elif lt == duration: # Offset last tick slightly left global_pos -= (duration * 0.03) custom_ticks.append(global_pos) custom_labels.append(str(lt)) ax1.set_xticks(custom_ticks) ax1.set_xticklabels(custom_labels, fontsize=8) # Save png_path = log_path.replace('.log', '.png') plt.tight_layout() plt.savefig(png_path) plt.close() try: rel_path = os.path.basename(png_path) if "output" in png_path: rel_path = f"output/{rel_path}" except: rel_path = png_path logger.info(f"Memory plot saved to: {rel_path}")