gen-worker 0.1.4__py3-none-any.whl

gen_worker/torch_manager/utils/base_types/architecture.py

@@ -0,0 +1,145 @@
+import torch
+from abc import ABC, abstractmethod
+from typing import Any, Iterable, TypeVar, Optional, TypedDict, Generic
+from spandrel import Architecture as SpandrelArchitecture, ImageModelDescriptor
+from .common import StateDict, TorchDevice
+
+T = TypeVar("T", bound=torch.nn.Module, covariant=True)
+
+
+class ComponentMetadata(TypedDict):
+    display_name: str
+    input_space: str
+    output_space: str
+
+
+# TO DO: in the future, maybe we can compare sets of keys, rather than use
+# a detect method? That might be more optimized.
+
+
+class Architecture(ABC, Generic[T]):
+    """
+    The abstract-base-class that all cozy-creator Architectures should implement.
+    """
+
+    @property
+    def display_name(self) -> str:
+        return self._display_name
+
+    @property
+    def input_space(self) -> str:
+        return self._input_space
+
+    @property
+    def output_space(self) -> str:
+        return self._output_space
+
+    @property
+    def model(self) -> T:
+        """Access the underlying PyTorch model."""
+        return self._model  # type: ignore
+
+    @property
+    def config(self) -> Any:
+        return self._config
+
+    def __init__(
+        self,
+        *,
+        state_dict: Optional[StateDict] = None,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """
+        Constructor signature should look like this, although this abstract-base
+        class does not (and cannot) enforce your constructor signature.
+        """
+        self._display_name = "default"
+        self._input_space = "default"
+        self._output_space = "default"
+        self._config = {}
+        pass
+
+    @classmethod
+    @abstractmethod
+    def detect(
+        cls,
+        *,
+        state_dict: Optional[StateDict] = None,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> Optional[ComponentMetadata]:
+        """
+        Detects whether the given state dictionary matches the architecture.
+
+        Args:
+            state_dict (StateDict): The state dictionary from a PyTorch model.
+            metadata (dict[str, Any]): optional additional metadata to help identify the model
+
+        Returns:
+            bool: True if the state dictionary matches the architecture, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def load(
+        self,
+        state_dict: StateDict,
+        device: Optional[TorchDevice] = None,
+    ) -> None:
+        """
+        Loads a model from the given state dictionary according to the architecture.
+
+        Args:
+            state_dict (StateDict): The state dictionary from a PyTorch model.
+            device: The device the loaded model is sent to.
+        """
+        pass
+
+
+
+class SpandrelArchitectureAdapter(Architecture):
+    """
+    This class converts architectures from the spandrel library to our own
+    Architecture interface.
+    """
+
+    def __init__(self, arch: SpandrelArchitecture):
+        super().__init__()
+        if not isinstance(arch, SpandrelArchitecture):
+            raise TypeError("'arch' must be an instance of spandrel Architecture")
+
+        self.inner = arch
+        self._model = None
+        self._display_name = self.inner.name
+
+    def load(self, state_dict: StateDict, device: Optional[TorchDevice] = None) -> None:
+        descriptor = self.inner.load(state_dict)
+        if not isinstance(descriptor, ImageModelDescriptor):
+            raise TypeError("descriptor must be an instance of ImageModelDescriptor")
+
+        self._model = descriptor.model
+        if device is not None:
+            self._model.to(device)
+        elif descriptor.supports_half:
+            self._model.to(torch.float16)
+        elif descriptor.supports_bfloat16:
+            self._model.to(torch.bfloat16)
+        else:
+            raise Exception("Device not provided and could not be inferred")
+
+    @classmethod
+    def detect(
+        cls,
+        state_dict: StateDict = None,
+        metadata: dict[str, Any] = None,
+    ) -> Optional[ComponentMetadata]:
+        pass
+
+
+def architecture_validator(plugin: Any) -> bool:
+    try:
+        if isinstance(plugin, Iterable):
+            return all(architecture_validator(p) for p in plugin)
+        return issubclass(plugin, Architecture)
+    except TypeError:
+        print(f"Invalid plugin type: {plugin}")
+        return False

gen_worker/torch_manager/utils/base_types/common.py

@@ -0,0 +1,52 @@
+import torch
+from typing import Union, List, Callable, Any
+import PIL.Image
+import numpy as np
+from enum import Enum
+from multiprocessing.connection import Connection
+
+
+class Language(Enum):
+    """
+    ISO 639-1 language codes; used for localizing text.
+    English will be displayed for all text lacking a localization.
+    """
+
+    ENGLISH = "en"
+    CHINESE = "zh"
+
+
+class Category(Enum):
+    """
+    Used to group nodes by category in the client.
+    """
+
+    LOADER = {Language.ENGLISH: "Loader", Language.CHINESE: "加载器"}
+    PIPE = {Language.ENGLISH: "Pipe", Language.CHINESE: "管道"}
+    UPSCALER = {Language.ENGLISH: "Upscaler", Language.CHINESE: "升频器"}
+    MASK = {Language.ENGLISH: "Mask"}
+    INPAINTING = {Language.ENGLISH: "Inpainting"}
+    IMAGES = {Language.ENGLISH: "Images"}
+
+
+StateDict = dict[str, torch.Tensor]
+"""
+The parameters of a PyTorch model, serialized as a flat dict.
+"""
+
+TorchDevice = Union[str, torch.device]
+"""
+A string like 'cpu', 'cuda', or 'mps' or a torch device object.
+"""
+
+ImageOutputType = Union[List[PIL.Image.Image], np.ndarray]
+"""
+Static typing for image outputs
+"""
+
+Validator = Callable[[Any], bool]
+
+JobQueueItem = tuple[dict[str, Any], Connection]
+"""
+Type of items on the job-queue
+"""

gen_worker/torch_manager/utils/base_types/config.py

@@ -0,0 +1,46 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional, Any
+import logging
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_HOME_DIR = os.path.expanduser("~/.cozy-creator/")
+
+
+@dataclass
+class PipelineConfig:
+    source: str
+    class_name: Optional[str | tuple[str, str]]
+    components: Optional[dict[str, "ComponentConfig"]]
+
+
+@dataclass
+class ComponentConfig:
+    source: str
+    class_name: Optional[str | tuple[str, str]]
+    kwargs: Optional[dict[str, Any]]
+
+
+def default_home_dir() -> str:
+    return DEFAULT_HOME_DIR
+
+
+def default_models_path() -> str:
+    return os.path.join(DEFAULT_HOME_DIR, "models")
+
+
+
+def default_host() -> str:
+    return "0.0.0.0" if os.path.exists("/.dockerenv") else "localhost"
+
+
+@dataclass
+class RuntimeConfig:
+    home_dir: str = field(default_factory=default_home_dir)
+    environment: str = "dev"
+    host: str = field(default_factory=default_host)
+    port: int = 8882
+    pipeline_defs: dict[str, PipelineConfig] = field(default_factory=dict)
+    enabled_models: list[str] = field(default_factory=list)
+    models_path: str = field(default_factory=default_models_path)

gen_worker/torch_manager/utils/config.py

@@ -0,0 +1,321 @@
+import os
+import argparse
+from typing import Optional, List, Callable, Dict, Any
+from .base_types.config import RuntimeConfig, PipelineConfig
+import yaml
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+cozy_config: Optional[RuntimeConfig] = None
+"""
+Global configuration for the Cozy Gen-Server
+"""
+
+
+def config_loaded() -> bool:
+    """
+    Returns a boolean indicating whether the config has been loaded.
+    This will return True if called within the cozy runtime, since the config is loaded at the start.
+    """
+    return cozy_config is not None
+
+
+def load_pipeline_defs_from_db(enabled_models: List[str]) -> Dict[str, Any]:
+    """
+    Load pipeline definitions from the database for the enabled models.
+    
+    Args:
+        enabled_models: List of model names to fetch from the database
+        
+    Returns:
+        Dictionary of pipeline definitions keyed by model name
+    """
+    try:
+        from .db.database import get_db_connection
+        from .repository import get_pipeline_defs
+        
+        # if not enabled_models:
+        #     return {}
+            
+        db_conn = get_db_connection()
+
+        if not db_conn:
+            logger.error("load_pipeline_defs_from_db: Could not get database connection.")
+            return {} # Cannot proceed without DB connection
+
+        names_to_query: List[str]
+        if enabled_models is None or not enabled_models: # Check if None or empty list
+            logger.info("load_pipeline_defs_from_db: No specific model names provided, fetching all model names from DB.")
+            all_db_model_names = []
+            with db_conn.cursor() as cursor: # Use a different variable name for cursor
+                cursor.execute("SELECT name FROM pipeline_defs WHERE source IS NOT NULL AND source != ''") # Fetch only usable models
+                rows = cursor.fetchall()
+                for row in rows:
+                    all_db_model_names.append(row['name'])
+            
+            if not all_db_model_names:
+                logger.warning("load_pipeline_defs_from_db: No model names found in DB to load definitions for.")
+                return {}
+            names_to_query = all_db_model_names
+        else:
+            names_to_query = enabled_models
+        
+        logger.debug(f"load_pipeline_defs_from_db: Fetching definitions for models: {names_to_query}")
+        
+        # get_pipeline_defs returns List[PipelineDef objects]
+        pipeline_def_objects = get_pipeline_defs(db_conn, names_to_query) 
+        
+        if not pipeline_def_objects:
+            logger.warning(f"load_pipeline_defs_from_db: get_pipeline_defs returned no objects for names: {names_to_query}")
+            return {}
+
+        logger.info(f"load_pipeline_defs_from_db: Loaded {len(pipeline_def_objects)} PipelineDef objects from repository.")
+        
+        # Convert DB models to dictionary format
+        db_pipeline_defs = {}
+        for model in pipeline_def_objects:
+            pipeline_def = {
+                "source": model.source,
+                "class_name": model.class_name,
+                "custom_pipeline": model.custom_pipeline,
+                "default_args": model.default_args,
+                "metadata": model.metadata,
+                "components": {}
+            }
+            
+            # Convert components
+            if model.components:
+                for name, comp in model.components.items():
+                    if isinstance(comp, dict):
+                        pipeline_def["components"][name] = {
+                            "class_name": comp.get("class_name", ""),
+                            "source": comp.get("source", ""),
+                            "kwargs": comp.get("kwargs", {})
+                        }
+            
+            # Add prompt definitions if available
+            if hasattr(model, "prompt_def") and model.prompt_def:
+                if not pipeline_def.get("default_args"):
+                    pipeline_def["default_args"] = {}
+                pipeline_def["default_args"]["positive_prompt"] = model.prompt_def.positive_prompt
+                pipeline_def["default_args"]["negative_prompt"] = model.prompt_def.negative_prompt
+                
+            db_pipeline_defs[model.name] = pipeline_def
+            
+        return db_pipeline_defs
+    except Exception as e:
+        logger.error(f"Error loading pipeline definitions from database: {e}")
+        return {}
+    
+
+def merge_pipeline_defs(existing_defs: Dict[str, Any], incoming_defs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Merge pipeline definitions from different sources, similar to Go implementation.
+    
+    Args:
+        existing_defs: Pipeline definitions from config.yaml
+        incoming_defs: Pipeline definitions from database
+        
+    Returns:
+        Merged pipeline definitions
+    """
+    merged_defs = existing_defs.copy()
+    
+    # Merge incoming defs into existing defs
+    for model_id, model_def in incoming_defs.items():
+        if model_id in merged_defs:
+            # Update only empty fields in existing definition
+            existing_def = merged_defs[model_id]
+            
+            if not existing_def.get("source"):
+                existing_def["source"] = model_def.get("source", "")
+                
+            if not existing_def.get("class_name"):
+                existing_def["class_name"] = model_def.get("class_name", "")
+                
+            if not existing_def.get("custom_pipeline"):
+                existing_def["custom_pipeline"] = model_def.get("custom_pipeline", "")
+                
+            if not existing_def.get("default_args"):
+                existing_def["default_args"] = model_def.get("default_args", {})
+                
+            if not existing_def.get("metadata"):
+                existing_def["metadata"] = model_def.get("metadata", {})
+                
+            # Handle components
+            if not existing_def.get("components"):
+                existing_def["components"] = {}
+                
+            # Merge components
+            for comp_name, comp_def in model_def.get("components", {}).items():
+                if comp_name in existing_def["components"]:
+                    # Update component fields if empty
+                    existing_comp = existing_def["components"][comp_name]
+                    
+                    if not existing_comp.get("class_name"):
+                        existing_comp["class_name"] = comp_def.get("class_name", "")
+                        
+                    if not existing_comp.get("source"):
+                        existing_comp["source"] = comp_def.get("source", "")
+                        
+                    if not existing_comp.get("kwargs"):
+                        existing_comp["kwargs"] = comp_def.get("kwargs", {})
+                else:
+                    # Add new component
+                    existing_def["components"][comp_name] = comp_def
+        else:
+            # Add new model definition
+            merged_defs[model_id] = model_def
+    
+    # Remove models without a source
+    models_to_remove = [model_id for model_id, def_obj in merged_defs.items() 
+                       if not def_obj.get("source")]
+    for model_id in models_to_remove:
+        del merged_defs[model_id]
+        
+    return merged_defs
+    
+
+def load_config() -> RuntimeConfig:
+    """
+    Load the configuration from a YAML file located at COZY_HOME/config.yaml.
+    Merges it with default values and database pipeline definitions.
+    """
+    default_home = os.path.expanduser("~/.cozy-creator")
+
+    cozy_mount_path = os.getenv("COZY_MOUNT_PATH")
+    if cozy_mount_path:
+        default_home = cozy_mount_path
+
+    default_models_path = os.path.join(default_home, "models")
+    
+    default_config = {
+        "home_dir": default_home,
+        "environment": "dev",
+        "host": "localhost",
+        "port": 8882,
+        "pipeline_defs": {},
+        # "enabled_models": [],
+        "models_path": default_models_path,
+    }
+
+    print(f"default_config: {default_config}")
+    
+    # Use COZY_HOME if set, else default.
+    home_dir = os.environ.get("COZY_HOME", default_home)
+    config_path = os.path.join(home_dir, "config.yaml")
+    
+    merged = default_config.copy()
+    
+    if os.path.exists(config_path):
+        try:
+            with open(config_path, "r") as f:
+                yaml_config = yaml.safe_load(f) or {}
+            
+            # Merge basic config values
+            # for key, value in yaml_config.items():
+            #     if key != "pipeline_defs":
+            #         merged[key] = value
+            
+            # Get pipeline defs from config
+            config_pipeline_defs = yaml_config.get("pipeline_defs", {})
+        except Exception as e:
+            logger.error(f"Error loading config from {config_path}: {e}")
+            config_pipeline_defs = {}
+    else:
+        logger.warning(f"Config file {config_path} not found. Using default configuration.")
+        config_pipeline_defs = {}
+    
+    # Get enabled models from environment if provided
+    # enabled_models_env = os.environ.get("ENABLED_MODELS")
+    # if enabled_models_env:
+    #     try:
+    #         # If provided as JSON.
+    #         merged["enabled_models"] = json.loads(enabled_models_env)
+    #     except Exception:
+    #         # Otherwise, assume comma-separated.
+    #         merged["enabled_models"] = [m.strip() for m in enabled_models_env.split(",") if m.strip()]
+    
+    # Load ALL pipeline definitions from the database
+    # The modified load_pipeline_defs_from_db with no args should fetch all.
+    logger.info("Loading all pipeline definitions from database...")
+    db_pipeline_defs = load_pipeline_defs_from_db(enabled_models=None)
+    print(f"db_pipeline_defs: {db_pipeline_defs}")
+    
+    # Merge pipeline definitions from config and database
+    merged["pipeline_defs"] = merge_pipeline_defs(config_pipeline_defs, db_pipeline_defs)
+    print(f"merged: {merged}")
+    
+    return RuntimeConfig(**merged)
+
+
+def set_config(config: RuntimeConfig):
+    """
+    Sets the global configuration object .
+    """
+    global cozy_config
+    cozy_config = config
+
+
+def set_environment(environment: str):
+    """
+    Sets the global environment variable.
+    """
+    global ENVIRONMENT
+    ENVIRONMENT = environment
+
+
+def get_environment() -> str:
+    """
+    Returns the global environment variable.
+    """
+    if ENVIRONMENT is None:
+        raise ValueError("Environment has not been set yet")
+
+    return ENVIRONMENT
+
+
+def get_config() -> RuntimeConfig:
+    """
+    Returns the global configuration object. This is only available if the config has been loaded, which happens at
+    the start of the server, else it will raise an error.
+    """
+    if cozy_config is None:
+        raise ValueError("Config has not been loaded yet")
+
+    return cozy_config
+
+
+ParseArgsMethod = Callable[
+    [argparse.ArgumentParser, Optional[List[str]], Optional[argparse.Namespace]],
+    Optional[argparse.Namespace],
+]
+
+
+def is_model_enabled(model_name: str) -> bool:
+    """
+    Returns a boolean indicating whether a model is enabled in the global configuration.
+    """
+    config = get_config()
+    if config.pipeline_defs is None:
+        return False
+
+    return model_name in config.pipeline_defs.keys()
+
+
+def get_mock_config() -> RuntimeConfig:
+    """
+    Returns a mock (or test) version of the global configuration object.
+    This can be used outside of the cozy server environment.
+    """
+
+    environment = "test"
+    # home_dir = DEFAULT_HOME_DIR
+
+    return RuntimeConfig(
+        port=8881,
+        host="127.0.0.1",
+        environment=environment,
+    )

gen_worker/torch_manager/utils/db/database.py

@@ -0,0 +1,46 @@
+import os
+import logging
+from typing import List, Optional
+import psycopg2
+from psycopg2.extras import RealDictCursor
+from dotenv import load_dotenv
+
+logger = logging.getLogger(__name__)
+
+# Load environment variables from .env file
+load_dotenv()
+
+_db_connection = None
+
+def get_db_connection():
+    """
+    Returns a connection to the database. Creates a new connection if one doesn't exist.
+    Uses the DB_DSN environment variable for connection details.
+    """
+    global _db_connection
+    
+    if _db_connection is None:
+        db_dsn = os.getenv("DB_DSN")
+        if not db_dsn:
+            raise ValueError("DB_DSN environment variable not set")
+        
+        try:
+            logger.info("Connecting to database...")
+            _db_connection = psycopg2.connect(db_dsn, cursor_factory=RealDictCursor)
+            logger.info("Database connection established")
+        except Exception as e:
+            logger.error(f"Error connecting to database: {e}")
+            raise
+    
+    return _db_connection
+
+def close_db_connection():
+    """
+    Closes the database connection if it exists.
+    """
+    global _db_connection
+    
+    if _db_connection is not None:
+        _db_connection.close()
+        _db_connection = None
+        logger.info("Database connection closed")

gen_worker/torch_manager/utils/device.py

@@ -0,0 +1,26 @@
+import logging
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def get_torch_device(index: int = 0) -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda", index)
+    if torch.backends.mps.is_available():
+        return torch.device("mps", index)
+    if torch.xpu.is_available():
+        return torch.device("xpu", index)
+
+    logger.warning("No device found, using CPU. This will slow down performance.")
+    return torch.device("cpu")
+
+
+def get_torch_device_count() -> int:
+    if torch.cuda.is_available():
+        return torch.cuda.device_count()
+    if torch.backends.mps.is_available():
+        return torch.mps.device_count()
+    if torch.xpu.is_available():
+        return torch.xpu.device_count()
+    return 1

gen_worker/torch_manager/utils/diffusers_fix.py

@@ -0,0 +1,10 @@
+def fix_sdxl_compat():
+    """Fix SDXL compatibility across diffusers versions"""
+    try:
+        from diffusers import StableDiffusionXLPipeline
+        if not hasattr(StableDiffusionXLPipeline, 'do_classifier_free_guidance'):
+            StableDiffusionXLPipeline.do_classifier_free_guidance = True
+    except:
+        pass
+
+fix_sdxl_compat()  # Auto-run on import