add core code

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
+```

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

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

core/pipeline_types/__init__.py

@@ -0,0 +1,4 @@
+from .base import BasePipelineCapture
+from .hunyuanvideo import HunyuanPipelineCapture
+
+__all__ = ["BasePipelineCapture", "HunyuanPipelineCapture"]

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

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

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"
+}

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