From 7b8a17a737ea27e2b1ca61bb3fff0a359d6f6bbd Mon Sep 17 00:00:00 2001 From: Fred Bliss Date: Mon, 20 Jan 2025 10:23:34 -0600 Subject: [PATCH] add core code --- README.md | 17 +++ core/analyzers.py | 99 +++++++++++++ core/data_store.py | 131 +++++++++++++++++ core/pipeline_types/__init__.py | 4 + core/pipeline_types/base.py | 27 ++++ core/pipeline_types/hunyuanvideo.py | 61 ++++++++ nodes.py | 212 ++++++++++++++++++++++++++++ requirements.txt | 6 + 8 files changed, 557 insertions(+) create mode 100644 core/analyzers.py create mode 100644 core/data_store.py create mode 100644 core/pipeline_types/__init__.py create mode 100644 core/pipeline_types/base.py create mode 100644 core/pipeline_types/hunyuanvideo.py diff --git a/README.md b/README.md index f783e60..8866b7c 100644 --- a/README.md +++ b/README.md @@ -25,3 +25,20 @@ As this will be maintained based on my own free time, the core will focus on the - Compare runs - Generate visualizations (UMAP, etc) - Produce structured reports + +## Tree + +``` +ComfyUI-EmbeddingPipelineAnalytics/ +├── __init__.py +├── nodes.py +├── core/ +│ ├── __init__.py +│ ├── data_store.py # core data storage/management +│ ├── analyzers.py # analysis implementations +│ └── pipeline_types/ # pipeline-specific implementations +│ ├── __init__.py +│ ├── base.py +│ └── hunyuan.py # hunyuanvideo-specific code +└── requirements.txt +``` diff --git a/core/analyzers.py b/core/analyzers.py new file mode 100644 index 0000000..9bb420f --- /dev/null +++ b/core/analyzers.py @@ -0,0 +1,99 @@ +from typing import Dict, Any, Optional, List +import torch +import numpy as np +from pathlib import Path +import umap +from sklearn.decomposition import PCA +from sklearn.manifold import TSNE + +class EmbeddingAnalyzer: + """Core analysis functionality for embeddings.""" + + @staticmethod + def calculate_statistics(tensor: torch.Tensor) -> Dict[str, float]: + """Calculate basic statistics for a tensor.""" + return { + "mean": float(tensor.mean()), + "std": float(tensor.std()), + "min": float(tensor.min()), + "max": float(tensor.max()), + "shape": list(tensor.shape) + } + + @staticmethod + def reduce_dimensions(tensor: torch.Tensor, + method: str = "umap", + n_components: int = 2) -> np.ndarray: + """Reduce dimensionality of embeddings.""" + # Flatten if needed + if tensor.ndim > 2: + tensor_2d = tensor.view(tensor.size(0), -1) + else: + tensor_2d = tensor + + tensor_np = tensor_2d.cpu().numpy() + + if method == "umap": + reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, n_components=n_components) + elif method == "pca": + reducer = PCA(n_components=n_components) + elif method == "tsne": + reducer = TSNE(n_components=n_components) + else: + raise ValueError(f"Unknown reduction method: {method}") + + return reducer.fit_transform(tensor_np) + + @staticmethod + def compare_embeddings(emb1: torch.Tensor, + emb2: torch.Tensor, + method: str = "cosine") -> float: + """Compare two sets of embeddings.""" + if emb1.ndim > 2: + emb1 = emb1.view(emb1.size(0), -1) + if emb2.ndim > 2: + emb2 = emb2.view(emb2.size(0), -1) + + if method == "cosine": + return torch.nn.functional.cosine_similarity( + emb1.mean(dim=0, keepdim=True), + emb2.mean(dim=0, keepdim=True) + ).item() + else: + raise ValueError(f"Unknown comparison method: {method}") + +class PipelineAnalyzer: + """Analyzes full pipeline runs.""" + + @staticmethod + def analyze_run(run_data: Dict[str, Any], + analysis_types: List[str] = ["statistics", "umap"]) -> Dict[str, Any]: + """Analyze a complete pipeline run.""" + analyzer = EmbeddingAnalyzer() + results = { + "run_id": run_data.get("run_id"), + "metadata": run_data.get("metadata", {}), + "analyses": {} + } + + # Analyze each stage's embeddings + for stage, embeddings in run_data.get("embeddings", {}).items(): + stage_results = {} + + for name, tensor in embeddings.items(): + tensor_results = {} + + if "statistics" in analysis_types: + tensor_results["statistics"] = analyzer.calculate_statistics(tensor) + + if "umap" in analysis_types: + tensor_results["umap"] = analyzer.reduce_dimensions(tensor, "umap").tolist() + + if "pca" in analysis_types: + tensor_results["pca"] = analyzer.reduce_dimensions(tensor, "pca").tolist() + + stage_results[name] = tensor_results + + results["analyses"][stage] = stage_results + + return results \ No newline at end of file diff --git a/core/data_store.py b/core/data_store.py new file mode 100644 index 0000000..ba49cfe --- /dev/null +++ b/core/data_store.py @@ -0,0 +1,131 @@ +# core/data_store.py +from typing import Dict, Any, Optional +from pathlib import Path +import json +import torch +import numpy as np +from datetime import datetime +import hashlib + +class PipelineDataCapture: + """Captures and stores pipeline execution data.""" + + def __init__(self, base_path: str = "pipeline_analytics"): + self.base_path = Path(base_path) + self.base_path.mkdir(parents=True, exist_ok=True) + + # Create subdirectories + self.embeddings_path = self.base_path / "embeddings" + self.metadata_path = self.base_path / "metadata" + self.outputs_path = self.base_path / "outputs" + self.analytics_path = self.base_path / "analytics" + + for path in [self.embeddings_path, self.metadata_path, + self.outputs_path, self.analytics_path]: + path.mkdir(exist_ok=True) + + def generate_run_id(self, metadata: Dict) -> str: + """Generate unique run ID based on timestamp and input hash.""" + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + input_hash = hashlib.sha256( + json.dumps(metadata, sort_keys=True).encode() + ).hexdigest()[:8] + return f"{timestamp}_{input_hash}" + + def save_embeddings(self, + run_id: str, + embeddings: Dict[str, torch.Tensor], + stage: str) -> Dict[str, str]: + """Save embeddings with metadata about stage in pipeline.""" + paths = {} + for name, tensor in embeddings.items(): + if tensor is None: + continue + + filename = f"{run_id}_{stage}_{name}.pt" + path = self.embeddings_path / filename + + # Save tensor + torch.save({ + "tensor": tensor, + "shape": tensor.shape, + "dtype": str(tensor.dtype), + "stage": stage, + "name": name + }, path) + + paths[name] = str(path) + + return paths + + def save_metadata(self, + run_id: str, + metadata: Dict[str, Any]) -> str: + """Save run metadata including parameters, prompt, etc.""" + filename = f"{run_id}_metadata.json" + path = self.metadata_path / filename + + with open(path, 'w') as f: + json.dump(metadata, f, indent=2) + + return str(path) + + def save_outputs(self, + run_id: str, + outputs: Dict[str, Any]) -> Dict[str, str]: + """Save pipeline outputs (images, videos, etc).""" + paths = {} + for name, output in outputs.items(): + if isinstance(output, torch.Tensor): + filename = f"{run_id}_{name}.pt" + path = self.outputs_path / filename + torch.save(output, path) + elif isinstance(output, np.ndarray): + filename = f"{run_id}_{name}.npy" + path = self.outputs_path / filename + np.save(path, output) + else: + filename = f"{run_id}_{name}.json" + path = self.outputs_path / filename + with open(path, 'w') as f: + json.dump(output, f) + + paths[name] = str(path) + + return paths + + def load_run(self, run_id: str) -> Dict[str, Any]: + """Load all data for a specific run.""" + run_data = { + "metadata": None, + "embeddings": {}, + "outputs": {}, + "analytics": {} + } + + # Load metadata + metadata_path = self.metadata_path / f"{run_id}_metadata.json" + if metadata_path.exists(): + with open(metadata_path) as f: + run_data["metadata"] = json.load(f) + + # Load embeddings + for path in self.embeddings_path.glob(f"{run_id}_*.pt"): + data = torch.load(path) + stage = data["stage"] + if stage not in run_data["embeddings"]: + run_data["embeddings"][stage] = {} + run_data["embeddings"][stage][data["name"]] = data["tensor"] + + # Load outputs + for path in self.outputs_path.glob(f"{run_id}_*"): + name = path.stem.replace(f"{run_id}_", "") + if path.suffix == ".pt": + run_data["outputs"][name] = torch.load(path) + elif path.suffix == ".npy": + run_data["outputs"][name] = np.load(path) + else: + with open(path) as f: + run_data["outputs"][name] = json.load(f) + + return run_data \ No newline at end of file diff --git a/core/pipeline_types/__init__.py b/core/pipeline_types/__init__.py new file mode 100644 index 0000000..78522ba --- /dev/null +++ b/core/pipeline_types/__init__.py @@ -0,0 +1,4 @@ +from .base import BasePipelineCapture +from .hunyuanvideo import HunyuanPipelineCapture + +__all__ = ["BasePipelineCapture", "HunyuanPipelineCapture"] \ No newline at end of file diff --git a/core/pipeline_types/base.py b/core/pipeline_types/base.py new file mode 100644 index 0000000..f7dec16 --- /dev/null +++ b/core/pipeline_types/base.py @@ -0,0 +1,27 @@ +from abc import ABC, abstractmethod +from typing import Dict, Any, Optional +from ..data_store import PipelineDataCapture + +class BasePipelineCapture(ABC): + """Base class for pipeline-specific capture implementations.""" + + def __init__(self, analytics_dir: str = "pipeline_analytics"): + self.data_store = PipelineDataCapture(analytics_dir) + + @abstractmethod + def capture_stage(self, + stage_name: str, + data: Any, + metadata: Optional[Dict] = None) -> str: + """Capture data from a pipeline stage.""" + pass + + @abstractmethod + def get_stage_types(self) -> Dict[str, type]: + """Return mapping of stage names to expected data types.""" + pass + + @abstractmethod + def validate_stage_data(self, stage_name: str, data: Any) -> bool: + """Validate data for a specific stage.""" + pass \ No newline at end of file diff --git a/core/pipeline_types/hunyuanvideo.py b/core/pipeline_types/hunyuanvideo.py new file mode 100644 index 0000000..ba188ac --- /dev/null +++ b/core/pipeline_types/hunyuanvideo.py @@ -0,0 +1,61 @@ +# core/pipeline_types/hunyuanvideo.py +from ..data_store import PipelineDataCapture +from typing import Dict, Any, Optional +import torch + +class HunyuanPipelineCapture: + """Specialized capture for HunyuanVideo pipeline.""" + + def __init__(self): + self.data_store = PipelineDataCapture("hunyuan_analytics") + + def capture_text_encode(self, + hyvid_embeds: Dict[str, torch.Tensor], + prompt: str, + template_used: Optional[str] = None, + cfg_settings: Optional[Dict] = None) -> str: + """Capture embeddings and metadata from text encoding stage.""" + + # Generate run ID from inputs + metadata = { + "stage": "text_encode", + "prompt": prompt, + "template": template_used, + "cfg_settings": cfg_settings + } + run_id = self.data_store.generate_run_id(metadata) + + # Save embeddings + embedding_paths = self.data_store.save_embeddings( + run_id, hyvid_embeds, "text_encode" + ) + metadata["embedding_paths"] = embedding_paths + + # Save metadata + self.data_store.save_metadata(run_id, metadata) + + return run_id + + def capture_sampling(self, + run_id: str, + samples: torch.Tensor, + sampling_params: Dict[str, Any]) -> str: + """Capture sampling parameters and outputs.""" + + # Load existing metadata + run_data = self.data_store.load_run(run_id) + metadata = run_data["metadata"] + + # Add sampling metadata + metadata["sampling_params"] = sampling_params + + # Save samples + output_paths = self.data_store.save_outputs( + run_id, {"samples": samples} + ) + metadata["output_paths"] = output_paths + + # Update metadata + self.data_store.save_metadata(run_id, metadata) + + return run_id \ No newline at end of file diff --git a/nodes.py b/nodes.py index e69de29..2cb5ffc 100644 --- a/nodes.py +++ b/nodes.py @@ -0,0 +1,212 @@ +import json +import torch +from typing import Dict, Any, Optional + +from .core.pipeline_types.hunyuan import HunyuanPipelineCapture + +class EmbeddingPipelineCapture: + """Base node for capturing pipeline data at any point.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "data": ("*",), # Accept any input type + "run_id": ("STRING", {"default": ""}), # Optional - will generate if empty + "stage_name": ("STRING", {}), + "metadata": ("STRING", {"multiline": True, "default": "{}"}) + } + } + + RETURN_TYPES = ("*", "STRING") # Pass through data + run_id + RETURN_NAMES = ("data", "run_id") + FUNCTION = "capture" + CATEGORY = "EmbeddingAnalytics" + + def __init__(self): + self.data_store = HunyuanPipelineCapture() # We can make this configurable later + + def capture(self, data: Any, run_id: str, stage_name: str, metadata: str) -> tuple: + try: + metadata_dict = json.loads(metadata) + except: + metadata_dict = {"raw": metadata} + + metadata_dict["stage"] = stage_name + + # Generate run_id if not provided + if not run_id: + run_id = self.data_store.data_store.generate_run_id(metadata_dict) + + # Handle different types of data + if isinstance(data, dict) and "samples" in data: # LATENT + self.data_store.data_store.save_outputs(run_id, {"latents": data["samples"]}) + elif isinstance(data, dict) and "prompt_embeds" in data: # HYVIDEMBEDS + self.data_store.data_store.save_embeddings(run_id, data, stage_name) + else: + # Save as generic output + self.data_store.data_store.save_outputs(run_id, {stage_name: data}) + + # Save/update metadata + self.data_store.data_store.save_metadata(run_id, metadata_dict) + + return (data, run_id) + +class EmbeddingAnalyzer: + """Node for analyzing captured embeddings.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "run_id": ("STRING", {}), + "analysis_type": (["umap", "pca", "tsne", "statistics"], {}), + }, + "optional": { + "compare_run_id": ("STRING", {"default": ""}), + "output_path": ("STRING", {"default": "embedding_analysis"}) + } + } + + RETURN_TYPES = ("STRING", "IMAGE") + RETURN_NAMES = ("analysis_results", "visualization") + FUNCTION = "analyze" + CATEGORY = "EmbeddingAnalytics" + + def __init__(self): + self.data_store = HunyuanPipelineCapture() + + def analyze(self, run_id: str, analysis_type: str, + compare_run_id: str = "", output_path: str = "embedding_analysis") -> tuple: + # Load run data + run_data = self.data_store.data_store.load_run(run_id) + + # Initialize analysis results + results = { + "run_id": run_id, + "analysis_type": analysis_type, + "stages": {} + } + + # Get comparison data if specified + compare_data = None + if compare_run_id: + compare_data = self.data_store.data_store.load_run(compare_run_id) + results["compare_run_id"] = compare_run_id + + # Analyze embeddings for each stage + for stage, embeddings in run_data["embeddings"].items(): + stage_results = analyze_embeddings( + embeddings, + analysis_type, + compare_embeddings=compare_data["embeddings"].get(stage) if compare_data else None + ) + results["stages"][stage] = stage_results + + # Generate visualization + visualization = generate_visualization( + results, + output_path=output_path + ) + + return (json.dumps(results, indent=2), visualization) + +def analyze_embeddings(embeddings: Dict[str, torch.Tensor], + analysis_type: str, + compare_embeddings: Optional[Dict[str, torch.Tensor]] = None) -> Dict: + """Analyze embeddings based on specified type.""" + results = {} + + for name, tensor in embeddings.items(): + # Basic statistics + stats = { + "mean": float(tensor.mean()), + "std": float(tensor.std()), + "min": float(tensor.min()), + "max": float(tensor.max()), + "shape": list(tensor.shape) + } + + # Dimensionality reduction if requested + if analysis_type in ["umap", "pca", "tsne"]: + # Prepare 2D tensor + if tensor.ndim > 2: + tensor_2d = tensor.view(tensor.size(0), -1) + else: + tensor_2d = tensor + + if analysis_type == "umap": + reducer = umap.UMAP(n_neighbors=15, min_dist=0.1) + reduced = reducer.fit_transform(tensor_2d.cpu().numpy()) + elif analysis_type == "pca": + reducer = PCA(n_components=2) + reduced = reducer.fit_transform(tensor_2d.cpu().numpy()) + else: # tsne + reducer = TSNE(n_components=2) + reduced = reducer.fit_transform(tensor_2d.cpu().numpy()) + + stats["reduced_coords"] = reduced.tolist() + + # Compare if provided + if compare_embeddings and name in compare_embeddings: + comp_tensor = compare_embeddings[name] + if comp_tensor.ndim > 2: + comp_tensor = comp_tensor.view(comp_tensor.size(0), -1) + + similarity = torch.nn.functional.cosine_similarity( + tensor_2d.mean(dim=0, keepdim=True), + comp_tensor.mean(dim=0, keepdim=True) + ).item() + stats["comparison_similarity"] = similarity + + results[name] = stats + + return results + +def generate_visualization(results: Dict, output_path: str) -> np.ndarray: + """Generate visualization of analysis results.""" + import matplotlib.pyplot as plt + from io import BytesIO + + plt.figure(figsize=(12, 8)) + + # Create subplots for each stage + num_stages = len(results["stages"]) + fig, axes = plt.subplots(num_stages, 1, figsize=(12, 6*num_stages)) + if num_stages == 1: + axes = [axes] + + for ax, (stage_name, stage_data) in zip(axes, results["stages"].items()): + # Plot embeddings if reduced coordinates available + for embed_name, embed_data in stage_data.items(): + if "reduced_coords" in embed_data: + coords = np.array(embed_data["reduced_coords"]) + ax.scatter(coords[:, 0], coords[:, 1], label=embed_name, alpha=0.6) + + ax.set_title(f"{stage_name} - {results['analysis_type'].upper()}") + ax.legend() + + plt.tight_layout() + + # Save plot + fig.savefig(f"{output_path}/analysis_{results['run_id']}.png") + + # Convert to image for ComfyUI + buf = BytesIO() + fig.savefig(buf, format='png') + buf.seek(0) + image = plt.imread(buf) + plt.close(fig) + + return image + +# Register nodes +NODE_CLASS_MAPPINGS = { + "EmbeddingPipelineCapture": EmbeddingPipelineCapture, + "EmbeddingAnalyzer": EmbeddingAnalyzer +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "EmbeddingPipelineCapture": "Embedding Pipeline Capture", + "EmbeddingAnalyzer": "Embedding Analyzer" +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index e69de29..38892f2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -0,0 +1,6 @@ +numpy>=1.21.0 +torch>=2.0.0 +umap-learn>=0.5.3 +scikit-learn>=1.0.0 +matplotlib>=3.4.0 +pandas>=1.3.0 \ No newline at end of file