diff --git a/skynet/config.py b/skynet/config.py index 21cc113..109c54c 100755 --- a/skynet/config.py +++ b/skynet/config.py @@ -3,8 +3,6 @@ import os import toml -from pathlib import Path - from .constants import DEFAULT_CONFIG_PATH diff --git a/skynet/constants.py b/skynet/constants.py index 6bc1867..480a5eb 100755 --- a/skynet/constants.py +++ b/skynet/constants.py @@ -1,22 +1,24 @@ #!/usr/bin/python +import msgspec +from typing import Literal VERSION = '0.1a12' DOCKER_RUNTIME_CUDA = 'skynet:runtime-cuda' -import msgspec -from typing import Literal class Size(msgspec.Struct): w: int h: int + class ModelDesc(msgspec.Struct): short: str mem: float size: Size tags: list[Literal['txt2img', 'img2img', 'inpaint']] + MODELS: dict[str, ModelDesc] = { 'runwayml/stable-diffusion-v1-5': ModelDesc( short='stable', diff --git a/skynet/frontend/discord/__init__.py b/skynet/frontend/discord/__init__.py old mode 100644 new mode 100755 index a6bc0f3..9bfcf08 --- a/skynet/frontend/discord/__init__.py +++ b/skynet/frontend/discord/__init__.py @@ -8,11 +8,18 @@ import asyncio from decimal import Decimal from hashlib import sha256 from datetime import datetime -from contextlib import ExitStack, AsyncExitStack +from contextlib import ( + ExitStack, + AsyncExitStack, +) from contextlib import asynccontextmanager as acm from leap.cleos import CLEOS -from leap.sugar import Name, asset_from_str, collect_stdout +from leap.sugar import ( + Name, + asset_from_str, + collect_stdout, +) from leap.hyperion import HyperionAPI # from telebot.types import InputMediaPhoto diff --git a/skynet/utils.py b/skynet/utils.py index 0662aca..24ac04f 100755 --- a/skynet/utils.py +++ b/skynet/utils.py @@ -9,9 +9,7 @@ import logging import importlib from typing import Optional -from pathlib import Path -import trio import torch import numpy as np @@ -112,6 +110,7 @@ def pipeline_for( return custom_pipeline.pipeline_for(model, mode, mem_fraction=mem_fraction, cache_dir=cache_dir) except ImportError: + # TODO, uhh why not warn/error log this? ... @@ -159,7 +158,7 @@ def pipeline_for( if mode == 'txt2img': pipe.vae.enable_tiling() pipe.vae.enable_slicing() - + pipe.enable_model_cpu_offload() else: