import torch import numpy as np import time import json import os from unittest.mock import patch import folder_paths try: # lpips分析图像差异距离 import lpips LPIPS_AVAILABLE = True except ImportError: LPIPS_AVAILABLE = False try: # 贝叶斯 from skopt import Optimizer from skopt.space import Real SKOPT_AVAILABLE = True except ImportError: SKOPT_AVAILABLE = False # 全局状态管理 class TeaCacheStateManager: _instance = None def __new__(cls): if cls._instance is None: cls._instance = super(TeaCacheStateManager, cls).__new__(cls) cls._instance.runs = {} cls._instance.lpips_model = None cls._instance.baseline_image = None return cls._instance def register_run(self, run_id, params): self.runs[run_id] = {"start_time": time.time(), "params": params, "results": {}, "cache": {}, "cnt": 0} def get_run_data(self, run_id): return self.runs.get(run_id) def cleanup_run(self, run_id): if run_id in self.runs: del self.runs[run_id] def set_lpips_model(self, model): self.lpips_model = model def get_lpips_model(self): return self.lpips_model def set_baseline_image(self, image_tensor): self.baseline_image = image_tensor def get_baseline_image(self): return self.baseline_image STATE_MANAGER = TeaCacheStateManager() # 缓存逻辑 def teacache_forward_analysis(self, x, timesteps, context, num_tokens, **kwargs): transformer_options = kwargs.get('transformer_options', {}) run_id = transformer_options.get('teacache_run_id') if not run_id: return self.forward_original(x, timesteps, context, num_tokens, **kwargs) run_data = STATE_MANAGER.get_run_data(run_id) if not run_data: return self.forward_original(x, timesteps, context, num_tokens, **kwargs) params = run_data['params'] if run_data.get("num_steps") is None and transformer_options.get("num_steps") is not None: run_data["num_steps"] = transformer_options.get("num_steps") cap_feats, cap_mask = context, kwargs.get('attention_mask') bs, c_channels, h_img, w_img = x.shape if hasattr(self, 'pad_to_patch_size'): x = self.pad_to_patch_size(x, (self.patch_size, self.patch_size)) else: ph = (self.patch_size - h_img % self.patch_size) % self.patch_size pw = (self.patch_size - w_img % self.patch_size) % self.patch_size x = torch.nn.functional.pad(x, (0, pw, 0, ph)) t = (1.0 - timesteps).to(dtype=x.dtype) t_emb = self.t_embedder(t, dtype=x.dtype) adaln_input = t_emb if cap_feats is not None: cap_feats = self.cap_embedder(cap_feats) x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t_emb, num_tokens) freqs_cis = freqs_cis.to(x.device) max_seq_len = x.shape[1] should_calc = True num_steps_in_state = run_data.get("num_steps") cnt = run_data.get('cnt', 0) if num_steps_in_state is None or num_steps_in_state == 0 or cnt == 0 or cnt == num_steps_in_state - 1: should_calc = True else: run_data.setdefault('cache', {}) current_cache = run_data['cache'].setdefault(max_seq_len, {"accumulated_rel_l1_distance": 0.0, "previous_modulated_input": None, "previous_residual": None}) if current_cache.get("previous_modulated_input") is not None: modulated_inp = self.layers[0].adaLN_modulation(adaln_input.clone())[0] coefficients = params.get("coefficients_to_use", []) rescale_func = np.poly1d(coefficients) prev_mod_input = current_cache["previous_modulated_input"] prev_mean = prev_mod_input.abs().mean().item() rel_l1_change = ((modulated_inp - prev_mod_input).abs().mean() / prev_mean).cpu().item() if prev_mean > 1e-9 else float('inf') current_cache["accumulated_rel_l1_distance"] += rescale_func(rel_l1_change) if current_cache["accumulated_rel_l1_distance"] < params.get('rel_l1_thresh', 0.3): should_calc = False else: should_calc = True current_cache["accumulated_rel_l1_distance"] = 0.0 current_cache["previous_modulated_input"] = modulated_inp.clone() else: should_calc = True if should_calc: if max_seq_len != run_data.get('uncond_seq_len'): run_data['results'].setdefault("total_inferences", 0) run_data['results']["total_inferences"] += 1 original_x = x.clone() processed_x = x for layer in self.layers: processed_x = layer(processed_x, mask, freqs_cis, adaln_input) run_data.setdefault('cache', {}) current_cache = run_data['cache'].setdefault(max_seq_len, {}) current_cache["previous_residual"] = processed_x - original_x current_cache["accumulated_rel_l1_distance"] = 0.0 if current_cache.get("previous_modulated_input") is None: current_cache["previous_modulated_input"] = self.layers[0].adaLN_modulation(adaln_input.clone())[0] else: current_cache = run_data['cache'].get(max_seq_len, {}) processed_x = x + current_cache.get("previous_residual", 0) if max_seq_len != run_data.get('uncond_seq_len'): run_data['results'].setdefault("cache_hits", 0) run_data['results']["cache_hits"] += 1 if num_steps_in_state is not None and max_seq_len != run_data.get('uncond_seq_len'): run_data['cnt'] += 1 output = self.final_layer(processed_x, adaln_input) if hasattr(self, 'unpatchify'): output = self.unpatchify(output, img_size, cap_size, return_tensor=True)[:, :, :h_img, :w_img] else: raise NotImplementedError("Model does not have an 'unpatchify' method.") return -output # TeaCache Patcher class TeaCache_Patcher: DEFAULT_COEFFS = [393.76566581, -603.50993606, 209.10239044, -23.00726601, 0.86377344] @classmethod def INPUT_TYPES(cls): modes = ["手动输入", "自动微调", "贝叶斯优化"] if SKOPT_AVAILABLE else ["手动输入", "自动微调"] metrics = ["速度-命中率权衡", "质量-命中率权衡 (LPIPS)"] if LPIPS_AVAILABLE else ["速度-命中率权衡"] optional_inputs = { "coefficients_str": ("STRING", {"multiline": True, "default": json.dumps(cls.DEFAULT_COEFFS)}), } if LPIPS_AVAILABLE: optional_inputs["max_lpips_thresh"] = ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.00001, "display": "number"}) return { "required": { "model": ("MODEL",), "mode": (modes,), "evaluation_metric": (metrics,), "rel_l1_thresh": ("FLOAT", {"default": 0.3, "min": 0.0, "step": 0.001}), }, "optional": optional_inputs } RETURN_TYPES = ("MODEL", "STRING") RETURN_NAMES = ("MODEL", "run_id") FUNCTION = "patch_model" CATEGORY = "utils/analysis" EXPERIMENTAL = True def _load_history(self, file_path): try: if os.path.exists(file_path): with open(file_path, 'r') as f: return json.load(f) except Exception: pass return [] def _get_score(self, run, baseline_time, metric, max_lpips_thresh): inferences = run.get("total_inferences", 0) if inferences == 0: return 0 hit_ratio = run.get("cache_hits", 0) / inferences if "LPIPS" in metric: lpips_dist = run.get("lpips_distance") if lpips_dist is None: return 0 if max_lpips_thresh > 0 and lpips_dist > max_lpips_thresh: return 0 if lpips_dist < 1e-6: return 0 # 公式：命中率 / LPIPS距离 return hit_ratio / lpips_dist else: # "速度-命中率权衡" time_saved = baseline_time - run.get('generation_time', 0) if time_saved <= 0: return 0 return time_saved * hit_ratio def _find_best_run(self, history, baseline_time, metric, max_lpips_thresh): eff_runs = [r for r in history if r.get("rel_l1_thresh") != 0] if not eff_runs: return None best_run, max_score = None, -1.0 for run in eff_runs: score = self._get_score(run, baseline_time, metric, max_lpips_thresh) if score > max_score: max_score, best_run = score, run return best_run def patch_model(self, model, mode, evaluation_metric, rel_l1_thresh, coefficients_str=None, max_lpips_thresh=0.6): history = self._load_history(os.path.join(folder_paths.get_output_directory(), "teacache_analysis.json")) baseline_runs = [r for r in history if r.get("rel_l1_thresh") == 0] # 基准时间只在“速度”相关的评估中需要 baseline_time = min([r['generation_time'] for r in baseline_runs]) if baseline_runs else None coeffs_to_use = [] if mode == "贝叶斯优化": if not SKOPT_AVAILABLE: raise Exception("'贝叶斯优化' 模式需要 'scikit-optimize' 库。") space = [Real(-1000, 1000), Real(-1000, 1000), Real(-1000, 1000), Real(-200, 200), Real(-50, 50)] optimizer = Optimizer(dimensions=space, random_state=int(time.time()), acq_func="gp_hedge") if history: valid_history = [r for r in history if "coefficients" in r and r.get("rel_l1_thresh") != 0 and len(r["coefficients"]) == len(space)] if "LPIPS" in evaluation_metric: valid_history = [r for r in valid_history if r.get("lpips_distance") is not None] elif "速度" in evaluation_metric and not baseline_time: print("警告: 缺少基准时间，无法进行基于速度的贝叶斯优化。") valid_history = [] if valid_history: y_iters = [-self._get_score(r, baseline_time, evaluation_metric, max_lpips_thresh) for r in valid_history] optimizer.tell([r["coefficients"] for r in valid_history], y_iters) coeffs_to_use = optimizer.ask() elif mode == "自动微调": best_efficient_run = self._find_best_run(history, baseline_time, evaluation_metric, max_lpips_thresh) base_coeffs = best_efficient_run["coefficients"] if best_efficient_run else self.DEFAULT_COEFFS idx_to_tweak = len(history) % len(base_coeffs) perturb_factor = np.random.uniform(0.8, 1.2) coeffs_to_use = list(base_coeffs) coeffs_to_use[idx_to_tweak] *= perturb_factor else: # 手动输入 coeffs_to_use = json.loads(coefficients_str) if coefficients_str else self.DEFAULT_COEFFS run_id = str(time.time_ns()) params_for_run = { "rel_l1_thresh": rel_l1_thresh, "coefficients_to_use": coeffs_to_use, "max_lpips_thresh": max_lpips_thresh } STATE_MANAGER.register_run(run_id, params_for_run) new_model = model.clone() diffusion_model = new_model.get_model_object("diffusion_model") if not hasattr(diffusion_model, 'forward_original'): diffusion_model.forward_original = diffusion_model.forward diffusion_model.forward = teacache_forward_analysis.__get__(diffusion_model, diffusion_model.__class__) def unet_wrapper_function(model_function, kwargs): c_dict = kwargs.get("c", {}) c_dict.setdefault("transformer_options", {}) c_dict["transformer_options"]["teacache_run_id"] = run_id if "sample_sigmas" in c_dict["transformer_options"]: c_dict["transformer_options"]["num_steps"] = len(c_dict["transformer_options"]["sample_sigmas"]) return model_function(kwargs["input"], kwargs["timestep"], **c_dict) new_model.set_model_unet_function_wrapper(unet_wrapper_function) print(f"[TeaCache Patcher] 已准备好运行，ID: {run_id}, 评估模式: {evaluation_metric}") return (new_model, run_id) # 结果收集 class TeaCache_Result_Collector: @classmethod def INPUT_TYPES(cls): return { "required": { "latent": ("LATENT",), "run_id": ("STRING", {"forceInput": True}), "analysis_file": ("STRING", {"default": "teacache_analysis.json"}), }, "optional": { "trigger": ("STRING", {"forceInput": True}) } } RETURN_TYPES = ("LATENT", "STRING") RETURN_NAMES = ("LATENT", "status") FUNCTION = "collect_and_save" CATEGORY = "utils/analysis" EXPERIMENTAL = True def _load_history(self, file_path): try: if os.path.exists(file_path): with open(file_path, 'r') as f: return json.load(f) except Exception: pass return [] def collect_and_save(self, latent, run_id, analysis_file, trigger=None): run_data = STATE_MANAGER.get_run_data(run_id) if not run_data: return (latent, f"错误: 未找到ID为 {run_id} 的运行记录。") params = run_data['params'] results = run_data['results'] generation_time = time.time() - run_data.get('start_time', time.time()) final_run_data = { "timestamp": time.time(), "generation_time": generation_time, "cache_hits": results.get('cache_hits', 0), "total_inferences": results.get('total_inferences', 0), "rel_l1_thresh": params.get('rel_l1_thresh'), "coefficients": params.get('coefficients_to_use'), "lpips_distance": results.get('lpips_distance', None), "max_lpips_thresh": params.get('max_lpips_thresh', None) } full_path = os.path.join(folder_paths.get_output_directory(), analysis_file) try: history_data = self._load_history(full_path) history_data.append(final_run_data) with open(full_path, 'w') as f: json.dump(history_data, f, indent=4) STATE_MANAGER.cleanup_run(run_id) summary = f"结果已成功保存到 {analysis_file}。\n" summary += f"耗时: {generation_time:.2f}s, 缓存命中/总数: {results.get('cache_hits', 0)}/{results.get('total_inferences', 0)}" if final_run_data["lpips_distance"] is not None: summary += f", LPIPS距离: {final_run_data['lpips_distance']:.4f}" print(summary) return (latent, summary) except Exception as e: return (latent, f"写入文件时出错: {e}") # LPIPS模型加载 class LPIPS_Model_Loader: def __init__(self): self.model = None @classmethod def INPUT_TYPES(cls): return {"required": {}} RETURN_TYPES = ("LPIPS_MODEL",) FUNCTION = "load_model" CATEGORY = "utils/analysis" EXPERIMENTAL = True def load_model(self): if not LPIPS_AVAILABLE: raise Exception("LPIPS库未安装。请执行 'pip install lpips'") if STATE_MANAGER.get_lpips_model() is None: print("正在加载LPIPS模型 (vgg)...") lpips_model = lpips.LPIPS(net='vgg').cpu() STATE_MANAGER.set_lpips_model(lpips_model) print("LPIPS模型加载完成。") return (STATE_MANAGER.get_lpips_model(),) # 基准图像存储 class Store_Baseline_Image: @classmethod def INPUT_TYPES(cls): return {"required": {"image": ("IMAGE",)}} RETURN_TYPES = ("BASELINE_IMG",) FUNCTION = "store_image" CATEGORY = "utils/analysis" EXPERIMENTAL = True def store_image(self, image): baseline_tensor = image.permute(0, 3, 1, 2).contiguous() STATE_MANAGER.set_baseline_image(baseline_tensor) return (baseline_tensor.clone(),) # LPIPS评估 class TeaCache_LPIPS_Evaluator: @classmethod def INPUT_TYPES(cls): return {"required": { "test_image": ("IMAGE",), "baseline_image": ("BASELINE_IMG",), "lpips_model": ("LPIPS_MODEL",), "run_id": ("STRING", {"forceInput": True}), }} RETURN_TYPES = ("STRING",) RETURN_NAMES = ("status",) FUNCTION = "evaluate" CATEGORY = "utils/analysis" EXPERIMENTAL = True def _preprocess_image(self, image_tensor): return image_tensor * 2.0 - 1.0 def evaluate(self, test_image, baseline_image, lpips_model, run_id): if baseline_image is None: return ("错误: 未提供基准图像 (Baseline Image)。",) run_data = STATE_MANAGER.get_run_data(run_id) if not run_data: return (f"错误: 未找到ID为 {run_id} 的运行记录。",) test_image_t = test_image.permute(0, 3, 1, 2) test_img_proc = self._preprocess_image(test_image_t) base_img_proc = self._preprocess_image(baseline_image) device = 'cpu' lpips_model.to(device) distance = lpips_model(test_img_proc.to(device), base_img_proc.to(device)) lpips_score = distance.item() run_data.setdefault('results', {})['lpips_distance'] = lpips_score return (f"LPIPS距离计算完成: {lpips_score:.4f}",)