skygpu
/
skynet
mirror of https://github.com/skygpu/skynet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #44 from skygpu/inpainting 

Add Inpainting to telos_testnet branch
pull/45/head
Guillermo Rodriguez 2025-01-20 20:54:17 -03:00 committed by GitHub
parent 7274fb017d 12d015f0b9
commit b5f52b3b5b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 3775 additions and 2707 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@ -29,9 +29,6 @@ poetry shell
# test you can run this command
skynet --help
# launch ipfs node
skynet run ipfs
# to launch worker
skynet run dgpu
@ -77,9 +74,6 @@ docker pull guilledk/skynet:runtime-cuda
# or build it (takes a bit of time)
./build_docker.sh
# launch simple ipfs node
./launch_ipfs.sh
# run worker with all gpus
docker run \
-it \

45
docker/Dockerfile.runtime+cuda-py312 100644
View File

@ -0,0 +1,45 @@
from nvidia/cuda:12.4.1-devel-ubuntu22.04
from python:3.12
env DEBIAN_FRONTEND=noninteractive
run apt-get update && apt-get install -y \
git \
llvm \
ffmpeg \
libsm6 \
libxext6 \
ninja-build
# env CC /usr/bin/clang
# env CXX /usr/bin/clang++
#
# # install llvm10 as required by llvm-lite
# run git clone https://github.com/llvm/llvm-project.git -b llvmorg-10.0.1
# workdir /llvm-project
# # this adds a commit from 12.0.0 that fixes build on newer compilers
# run git cherry-pick -n b498303066a63a203d24f739b2d2e0e56dca70d1
# run cmake -S llvm -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
# run ninja -C build install # -j8
run curl -sSL https://install.python-poetry.org | python3 -
env PATH "/root/.local/bin:$PATH"
copy . /skynet
workdir /skynet
env POETRY_VIRTUALENVS_PATH /skynet/.venv
run poetry install --with=cuda -v
workdir /root/target
env PYTORCH_CUDA_ALLOC_CONF max_split_size_mb:128
env NVIDIA_VISIBLE_DEVICES=all
copy docker/entrypoint.sh /entrypoint.sh
entrypoint ["/entrypoint.sh"]
cmd ["skynet", "--help"]

23
docker/build_docker.sh
View File

@ -1,20 +1,7 @@
docker build \
-t guilledk/skynet:runtime \
-f docker/Dockerfile.runtime .
-t guilledk/skynet:runtime-cuda-py312 \
-f docker/Dockerfile.runtime+cuda-py312 .
docker build \
-t guilledk/skynet:runtime-frontend \
-f docker/Dockerfile.runtime+frontend .
docker build \
-t guilledk/skynet:runtime-cuda-py311 \
-f docker/Dockerfile.runtime+cuda-py311 .
docker build \
-t guilledk/skynet:runtime-cuda \
-f docker/Dockerfile.runtime+cuda-py311 .
docker build \
-t guilledk/skynet:runtime-cuda-py310 \
-f docker/Dockerfile.runtime+cuda-py310 .
# docker build \
# -t guilledk/skynet:runtime-cuda \
# -f docker/Dockerfile.runtime+cuda-py311 .

4977
poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

45
pyproject.toml
View File

@ -1,20 +1,25 @@
[tool.poetry]
name = 'skynet'
version = '0.1a12'
version = '0.1a13'
description = 'Decentralized compute platform'
authors = ['Guillermo Rodriguez <guillermo@telos.net>']
license = 'AGPL'
readme = 'README.md'
[tool.poetry.dependencies]
python = '>=3.10,<3.12'
python = '>=3.10,<3.13'
pytz = '^2023.3.post1'
trio = '^0.22.2'
asks = '^3.0.0'
Pillow = '^10.0.1'
docker = '^6.1.3'
py-leap = {git = 'https://github.com/guilledk/py-leap.git', rev = 'v0.1a14'}
py-leap = {git = 'https://github.com/guilledk/py-leap.git', rev = 'v0.1a32'}
toml = '^0.10.2'
msgspec = "^0.19.0"
numpy = "<2.1"
protobuf = "^5.29.3"
zstandard = "^0.23.0"
click = "^8.1.8"
httpx = "^0.28.1"
[tool.poetry.group.frontend]
optional = true
@ -32,31 +37,33 @@ optional = true
[tool.poetry.group.dev.dependencies]
pdbpp = {version = '^0.10.3'}
pytest = {version = '^7.4.2'}
pytest-trio = "^0.8.0"
[tool.poetry.group.cuda]
optional = true
[tool.poetry.group.cuda.dependencies]
torch = {version = '2.0.1+cu118', source = 'torch'}
scipy = {version = '^1.11.2'}
numba = {version = '0.57.0'}
torch = {version = '2.5.1+cu121', source = 'torch'}
scipy = {version = '1.15.1'}
numba = {version = '0.60.0'}
quart = {version = '^0.19.3'}
triton = {version = '2.0.0', source = 'torch'}
basicsr = {version = '^1.4.2'}
xformers = {version = '^0.0.22'}
triton = {version = '3.1.0', source = 'torch'}
xformers = {version = '^0.0.29'}
hypercorn = {version = '^0.14.4'}
diffusers = {version = '^0.21.2'}
realesrgan = {version = '^0.3.0'}
diffusers = {version = '0.32.1'}
quart-trio = {version = '^0.11.0'}
torchvision = {version = '0.15.2+cu118', source = 'torch'}
accelerate = {version = '^0.23.0'}
transformers = {version = '^4.33.2'}
huggingface-hub = {version = '^0.17.3'}
torchvision = {version = '0.20.1+cu121', source = 'torch'}
accelerate = {version = '0.34.0'}
transformers = {version = '4.48.0'}
huggingface-hub = {version = '^0.27.1'}
invisible-watermark = {version = '^0.2.0'}
bitsandbytes = "^0.45.0"
basicsr = "^1.4.2"
realesrgan = "^0.3.0"
[[tool.poetry.source]]
name = 'torch'
url = 'https://download.pytorch.org/whl/cu118'
url = 'https://download.pytorch.org/whl/cu121'
priority = 'explicit'
[build-system]
@ -65,3 +72,7 @@ build-backend = 'poetry.core.masonry.api'
[tool.poetry.scripts]
skynet = 'skynet.cli:skynet'
txt2img = 'skynet.cli:txt2img'
img2img = 'skynet.cli:img2img'
upscale = 'skynet.cli:upscale'
inpaint = 'skynet.cli:inpaint'

39
skynet/cli.py
View File

@ -8,7 +8,7 @@ from functools import partial
import click
from leap.sugar import Name, asset_from_str
from leap.protocol import Name, Asset
from .config import *
from .constants import *
@ -20,7 +20,7 @@ def skynet(*args, **kwargs):
@click.command()
@click.option('--model', '-m', default='midj')
@click.option('--model', '-m', default=list(MODELS.keys())[-1])
@click.option(
'--prompt', '-p', default='a red old tractor in a sunny wheat field')
@click.option('--output', '-o', default='output.png')
@ -39,7 +39,7 @@ def txt2img(*args, **kwargs):
utils.txt2img(hf_token, **kwargs)
@click.command()
@click.option('--model', '-m', default=list(MODELS.keys())[0])
@click.option('--model', '-m', default=list(MODELS.keys())[-2])
@click.option(
'--prompt', '-p', default='a red old tractor in a sunny wheat field')
@click.option('--input', '-i', default='input.png')
@ -66,6 +66,37 @@ def img2img(model, prompt, input, output, strength, guidance, steps, seed):
seed=seed
)
@click.command()
@click.option('--model', '-m', default=list(MODELS.keys())[-3])
@click.option(
'--prompt', '-p', default='a red old tractor in a sunny wheat field')
@click.option('--input', '-i', default='input.png')
@click.option('--mask', '-M', default='mask.png')
@click.option('--output', '-o', default='output.png')
@click.option('--strength', '-Z', default=1.0)
@click.option('--guidance', '-g', default=10.0)
@click.option('--steps', '-s', default=26)
@click.option('--seed', '-S', default=None)
def inpaint(model, prompt, input, mask, output, strength, guidance, steps, seed):
from . import utils
config = load_skynet_toml()
hf_token = load_key(config, 'skynet.dgpu.hf_token')
hf_home = load_key(config, 'skynet.dgpu.hf_home')
set_hf_vars(hf_token, hf_home)
utils.inpaint(
hf_token,
model=model,
prompt=prompt,
img_path=input,
mask_path=mask,
output=output,
strength=strength,
guidance=guidance,
steps=steps,
seed=seed
)
@click.command()
@click.option('--input', '-i', default='input.png')
@click.option('--output', '-o', default='output.png')
@ -147,7 +178,7 @@ def enqueue(
'user': Name(account),
'request_body': req,
'binary_data': binary,
'reward': asset_from_str(reward),
'reward': Asset.from_str(reward),
'min_verification': 1
},
account, key, permission,

343
skynet/constants.py
View File

@ -4,31 +4,120 @@ VERSION = '0.1a12'
DOCKER_RUNTIME_CUDA = 'skynet:runtime-cuda'
MODELS = {
'prompthero/openjourney': {'short': 'midj', 'mem': 6},
'runwayml/stable-diffusion-v1-5': {'short': 'stable', 'mem': 6},
'stabilityai/stable-diffusion-2-1-base': {'short': 'stable2', 'mem': 6},
'snowkidy/stable-diffusion-xl-base-0.9': {'short': 'stablexl0.9', 'mem': 8.3},
'Linaqruf/anything-v3.0': {'short': 'hdanime', 'mem': 6},
'hakurei/waifu-diffusion': {'short': 'waifu', 'mem': 6},
'nitrosocke/Ghibli-Diffusion': {'short': 'ghibli', 'mem': 6},
'dallinmackay/Van-Gogh-diffusion': {'short': 'van-gogh', 'mem': 6},
'lambdalabs/sd-pokemon-diffusers': {'short': 'pokemon', 'mem': 6},
'Envvi/Inkpunk-Diffusion': {'short': 'ink', 'mem': 6},
'nousr/robo-diffusion': {'short': 'robot', 'mem': 6},
import msgspec
from typing import Literal
# default is always last
'stabilityai/stable-diffusion-xl-base-1.0': {'short': 'stablexl', 'mem': 8.3},
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',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'stabilityai/stable-diffusion-2-1-base': ModelDesc(
short='stable2',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'snowkidy/stable-diffusion-xl-base-0.9': ModelDesc(
short='stablexl0.9',
mem=8.3,
size=Size(w=1024, h=1024),
tags=['txt2img']
),
'Linaqruf/anything-v3.0': ModelDesc(
short='hdanime',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'hakurei/waifu-diffusion': ModelDesc(
short='waifu',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'nitrosocke/Ghibli-Diffusion': ModelDesc(
short='ghibli',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'dallinmackay/Van-Gogh-diffusion': ModelDesc(
short='van-gogh',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'lambdalabs/sd-pokemon-diffusers': ModelDesc(
short='pokemon',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'Envvi/Inkpunk-Diffusion': ModelDesc(
short='ink',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'nousr/robo-diffusion': ModelDesc(
short='robot',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img']
),
'black-forest-labs/FLUX.1-schnell': ModelDesc(
short='flux',
mem=24,
size=Size(w=1024, h=1024),
tags=['txt2img']
),
'black-forest-labs/FLUX.1-Fill-dev': ModelDesc(
short='flux-inpaint',
mem=24,
size=Size(w=1024, h=1024),
tags=['inpaint']
),
'diffusers/stable-diffusion-xl-1.0-inpainting-0.1': ModelDesc(
short='stablexl-inpaint',
mem=8.3,
size=Size(w=1024, h=1024),
tags=['inpaint']
),
'prompthero/openjourney': ModelDesc(
short='midj',
mem=6,
size=Size(w=512, h=512),
tags=['txt2img', 'img2img']
),
'stabilityai/stable-diffusion-xl-base-1.0': ModelDesc(
short='stablexl',
mem=8.3,
size=Size(w=1024, h=1024),
tags=['txt2img']
),
}
SHORT_NAMES = [
model_info['short']
model_info.short
for model_info in MODELS.values()
]
def get_model_by_shortname(short: str):
for model, info in MODELS.items():
if short == info['short']:
if short == info.short:
return model
N = '\n'
@ -166,9 +255,7 @@ DEFAULT_UPSCALER = None
DEFAULT_CONFIG_PATH = 'skynet.toml'
DEFAULT_INITAL_MODELS = [
'stabilityai/stable-diffusion-xl-base-1.0'
]
DEFAULT_INITAL_MODEL = list(MODELS.keys())[-1]
DATE_FORMAT = '%B the %dth %Y, %H:%M:%S'
@ -193,3 +280,221 @@ TG_MAX_WIDTH = 1280
TG_MAX_HEIGHT = 1280
DEFAULT_SINGLE_CARD_MAP = 'cuda:0'
GPU_CONTRACT_ABI = {
"version": "eosio::abi/1.2",
"types": [],
"structs": [
{
"name": "account",
"base": "",
"fields": [
{"name": "user", "type": "name"},
{"name": "balance", "type": "asset"},
{"name": "nonce", "type": "uint64"}
]
},
{
"name": "card",
"base": "",
"fields": [
{"name": "id", "type": "uint64"},
{"name": "owner", "type": "name"},
{"name": "card_name", "type": "string"},
{"name": "version", "type": "string"},
{"name": "total_memory", "type": "uint64"},
{"name": "mp_count", "type": "uint32"},
{"name": "extra", "type": "string"}
]
},
{
"name": "clean",
"base": "",
"fields": []
},
{
"name": "config",
"base": "",
"fields": [
{"name": "token_contract", "type": "name"},
{"name": "token_symbol", "type": "symbol"}
]
},
{
"name": "dequeue",
"base": "",
"fields": [
{"name": "user", "type": "name"},
{"name": "request_id", "type": "uint64"}
]
},
{
"name": "enqueue",
"base": "",
"fields": [
{"name": "user", "type": "name"},
{"name": "request_body", "type": "string"},
{"name": "binary_data", "type": "string"},
{"name": "reward", "type": "asset"},
{"name": "min_verification", "type": "uint32"}
]
},
{
"name": "gcfgstruct",
"base": "",
"fields": [
{"name": "token_contract", "type": "name"},
{"name": "token_symbol", "type": "symbol"}
]
},
{
"name": "submit",
"base": "",
"fields": [
{"name": "worker", "type": "name"},
{"name": "request_id", "type": "uint64"},
{"name": "request_hash", "type": "checksum256"},
{"name": "result_hash", "type": "checksum256"},
{"name": "ipfs_hash", "type": "string"}
]
},
{
"name": "withdraw",
"base": "",
"fields": [
{"name": "user", "type": "name"},
{"name": "quantity", "type": "asset"}
]
},
{
"name": "work_request_struct",
"base": "",
"fields": [
{"name": "id", "type": "uint64"},
{"name": "user", "type": "name"},
{"name": "reward", "type": "asset"},
{"name": "min_verification", "type": "uint32"},
{"name": "nonce", "type": "uint64"},
{"name": "body", "type": "string"},
{"name": "binary_data", "type": "string"},
{"name": "timestamp", "type": "time_point_sec"}
]
},
{
"name": "work_result_struct",
"base": "",
"fields": [
{"name": "id", "type": "uint64"},
{"name": "request_id", "type": "uint64"},
{"name": "user", "type": "name"},
{"name": "worker", "type": "name"},
{"name": "result_hash", "type": "checksum256"},
{"name": "ipfs_hash", "type": "string"},
{"name": "submited", "type": "time_point_sec"}
]
},
{
"name": "workbegin",
"base": "",
"fields": [
{"name": "worker", "type": "name"},
{"name": "request_id", "type": "uint64"},
{"name": "max_workers", "type": "uint32"}
]
},
{
"name": "workcancel",
"base": "",
"fields": [
{"name": "worker", "type": "name"},
{"name": "request_id", "type": "uint64"},
{"name": "reason", "type": "string"}
]
},
{
"name": "worker",
"base": "",
"fields": [
{"name": "account", "type": "name"},
{"name": "joined", "type": "time_point_sec"},
{"name": "left", "type": "time_point_sec"},
{"name": "url", "type": "string"}
]
},
{
"name": "worker_status_struct",
"base": "",
"fields": [
{"name": "worker", "type": "name"},
{"name": "status", "type": "string"},
{"name": "started", "type": "time_point_sec"}
]
}
],
"actions": [
{"name": "clean", "type": "clean", "ricardian_contract": ""},
{"name": "config", "type": "config", "ricardian_contract": ""},
{"name": "dequeue", "type": "dequeue", "ricardian_contract": ""},
{"name": "enqueue", "type": "enqueue", "ricardian_contract": ""},
{"name": "submit", "type": "submit", "ricardian_contract": ""},
{"name": "withdraw", "type": "withdraw", "ricardian_contract": ""},
{"name": "workbegin", "type": "workbegin", "ricardian_contract": ""},
{"name": "workcancel", "type": "workcancel", "ricardian_contract": ""}
],
"tables": [
{
"name": "cards",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "card"
},
{
"name": "gcfgstruct",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "gcfgstruct"
},
{
"name": "queue",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "work_request_struct"
},
{
"name": "results",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "work_result_struct"
},
{
"name": "status",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "worker_status_struct"
},
{
"name": "users",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "account"
},
{
"name": "workers",
"index_type": "i64",
"key_names": [],
"key_types": [],
"type": "worker"
}
],
"ricardian_clauses": [],
"error_messages": [],
"abi_extensions": [],
"variants": [],
"action_results": []
}

187
skynet/dgpu/compute.py
View File

@ -13,36 +13,48 @@ from diffusers import DiffusionPipeline
import trio
import torch
from skynet.constants import DEFAULT_INITAL_MODELS, MODELS
from skynet.constants import DEFAULT_INITAL_MODEL, MODELS
from skynet.dgpu.errors import DGPUComputeError, DGPUInferenceCancelled
from skynet.utils import crop_image, convert_from_cv2_to_image, convert_from_image_to_cv2, convert_from_img_to_bytes, init_upscaler, pipeline_for
def prepare_params_for_diffuse(
params: dict,
input_type: str,
binary = None
mode: str,
inputs: list[bytes]
):
_params = {}
if binary != None:
match input_type:
case 'png':
image = crop_image(
binary, params['width'], params['height'])
match mode:
case 'inpaint':
image = crop_image(
inputs[0], params['width'], params['height'])
_params['image'] = image
mask = crop_image(
inputs[1], params['width'], params['height'])
_params['image'] = image
_params['mask_image'] = mask
if 'flux' in params['model'].lower():
_params['max_sequence_length'] = 512
else:
_params['strength'] = float(params['strength'])
case 'none':
...
case 'img2img':
image = crop_image(
inputs[0], params['width'], params['height'])
case _:
raise DGPUComputeError(f'Unknown input_type {input_type}')
_params['image'] = image
_params['strength'] = float(params['strength'])
else:
_params['width'] = int(params['width'])
_params['height'] = int(params['height'])
case 'txt2img' | 'diffuse':
...
case _:
raise DGPUComputeError(f'Unknown mode {mode}')
# _params['width'] = int(params['width'])
# _params['height'] = int(params['height'])
return (
params['prompt'],
@ -57,95 +69,55 @@ def prepare_params_for_diffuse(
class SkynetMM:
def __init__(self, config: dict):
self.upscaler = init_upscaler()
self.initial_models = (
config['initial_models']
if 'initial_models' in config else DEFAULT_INITAL_MODELS
)
self.cache_dir = None
if 'hf_home' in config:
self.cache_dir = config['hf_home']
self._models = {}
for model in self.initial_models:
self.load_model(model, False, force=True)
self._model_name = ''
self._model_mode = ''
# self.load_model(DEFAULT_INITAL_MODEL, 'txt2img')
def log_debug_info(self):
logging.info('memory summary:')
logging.info('\n' + torch.cuda.memory_summary())
def is_model_loaded(self, model_name: str, image: bool):
for model_key, model_data in self._models.items():
if (model_key == model_name and
model_data['image'] == image):
return True
def is_model_loaded(self, name: str, mode: str):
if (name == self._model_name and
mode == self._model_mode):
return True
return False
def unload_model(self):
if getattr(self, '_model', None):
del self._model
gc.collect()
torch.cuda.empty_cache()
self._model_name = ''
self._model_mode = ''
def load_model(
self,
model_name: str,
image: bool,
force=False
name: str,
mode: str
):
logging.info(f'loading model {model_name}...')
if force or len(self._models.keys()) == 0:
pipe = pipeline_for(
model_name, image=image, cache_dir=self.cache_dir)
logging.info(f'loading model {name}...')
self.unload_model()
self._model = pipeline_for(
name, mode, cache_dir=self.cache_dir)
self._model_mode = mode
self._model_name = name
self._models[model_name] = {
'pipe': pipe,
'generated': 0,
'image': image
}
else:
least_used = list(self._models.keys())[0]
for model in self._models:
if self._models[
least_used]['generated'] > self._models[model]['generated']:
least_used = model
del self._models[least_used]
logging.info(f'swapping model {least_used} for {model_name}...')
gc.collect()
torch.cuda.empty_cache()
pipe = pipeline_for(
model_name, image=image, cache_dir=self.cache_dir)
self._models[model_name] = {
'pipe': pipe,
'generated': 0,
'image': image
}
logging.info(f'loaded model {model_name}')
return pipe
def get_model(self, model_name: str, image: bool) -> DiffusionPipeline:
if model_name not in MODELS:
raise DGPUComputeError(f'Unknown model {model_name}')
if not self.is_model_loaded(model_name, image):
pipe = self.load_model(model_name, image=image)
else:
pipe = self._models[model_name]['pipe']
return pipe
def compute_one(
self,
request_id: int,
method: str,
params: dict,
input_type: str = 'png',
binary: bytes | None = None
inputs: list[bytes] = []
):
def maybe_cancel_work(step, *args, **kwargs):
if self._should_cancel:
@ -154,6 +126,8 @@ class SkynetMM:
logging.warn(f'cancelling work at step {step}')
raise DGPUInferenceCancelled()
return {}
maybe_cancel_work(0)
output_type = 'png'
@ -163,20 +137,29 @@ class SkynetMM:
output = None
output_hash = None
try:
match method:
case 'diffuse':
arguments = prepare_params_for_diffuse(
params, input_type, binary=binary)
prompt, guidance, step, seed, upscaler, extra_params = arguments
model = self.get_model(params['model'], 'image' in extra_params)
name = params['model']
output = model(
match method:
case 'diffuse' | 'txt2img' | 'img2img' | 'inpaint':
if not self.is_model_loaded(name, method):
self.load_model(name, method)
arguments = prepare_params_for_diffuse(
params, method, inputs)
prompt, guidance, step, seed, upscaler, extra_params = arguments
if 'flux' in name.lower():
extra_params['callback_on_step_end'] = maybe_cancel_work
else:
extra_params['callback'] = maybe_cancel_work
extra_params['callback_steps'] = 1
output = self._model(
prompt,
guidance_scale=guidance,
num_inference_steps=step,
generator=seed,
callback=maybe_cancel_work,
callback_steps=1,
**extra_params
).images[0]
@ -185,7 +168,7 @@ class SkynetMM:
case 'png':
if upscaler == 'x4':
input_img = output.convert('RGB')
up_img, _ = self.upscaler.enhance(
up_img, _ = init_upscaler().enhance(
convert_from_image_to_cv2(input_img), outscale=4)
output = convert_from_cv2_to_image(up_img)
@ -197,6 +180,22 @@ class SkynetMM:
output_hash = sha256(output_binary).hexdigest()
case 'upscale':
if self._model_mode != 'upscale':
self.unload_model()
self._model = init_upscaler()
self._model_mode = 'upscale'
self._model_name = 'realesrgan'
input_img = inputs[0].convert('RGB')
up_img, _ = self._model.enhance(
convert_from_image_to_cv2(input_img), outscale=4)
output = convert_from_cv2_to_image(up_img)
output_binary = convert_from_img_to_bytes(output)
output_hash = sha256(output_binary).hexdigest()
case _:
raise DGPUComputeError('Unsupported compute method')

185
skynet/dgpu/daemon.py
View File

@ -117,6 +117,103 @@ class SkynetDGPUDaemon:
return app
async def maybe_serve_one(self, req):
rid = req['id']
# parse request
body = json.loads(req['body'])
model = body['params']['model']
# if model not known
if model != 'RealESRGAN_x4plus' and model not in MODELS:
logging.warning(f'Unknown model {model}')
return False
# if whitelist enabled and model not in it continue
if (len(self.model_whitelist) > 0 and
not model in self.model_whitelist):
return False
# if blacklist contains model skip
if model in self.model_blacklist:
return False
my_results = [res['id'] for res in self._snap['my_results']]
if rid not in my_results and rid in self._snap['requests']:
statuses = self._snap['requests'][rid]
if len(statuses) == 0:
inputs = []
for _input in req['binary_data'].split(','):
if _input:
for _ in range(3):
try:
img = await self.conn.get_input_data(_input)
inputs.append(img)
break
except:
...
hash_str = (
str(req['nonce'])
+
req['body']
+
req['binary_data']
)
logging.info(f'hashing: {hash_str}')
request_hash = sha256(hash_str.encode('utf-8')).hexdigest()
# TODO: validate request
# perform work
logging.info(f'working on {body}')
resp = await self.conn.begin_work(rid)
if not resp or 'code' in resp:
logging.info(f'probably being worked on already... skip.')
else:
try:
output_type = 'png'
if 'output_type' in body['params']:
output_type = body['params']['output_type']
output = None
output_hash = None
match self.backend:
case 'sync-on-thread':
self.mm._should_cancel = self.should_cancel_work
output_hash, output = await trio.to_thread.run_sync(
partial(
self.mm.compute_one,
rid,
body['method'], body['params'],
inputs=inputs
)
)
case _:
raise DGPUComputeError(f'Unsupported backend {self.backend}')
self._last_generation_ts = datetime.now().isoformat()
self._last_benchmark = self._benchmark
self._benchmark = []
ipfs_hash = await self.conn.publish_on_ipfs(output, typ=output_type)
await self.conn.submit_work(rid, request_hash, output_hash, ipfs_hash)
except BaseException as e:
traceback.print_exc()
await self.conn.cancel_work(rid, str(e))
finally:
return True
else:
logging.info(f'request {rid} already beign worked on, skip...')
async def serve_forever(self):
try:
while True:
@ -133,92 +230,8 @@ class SkynetDGPUDaemon:
)
for req in queue:
rid = req['id']
# parse request
body = json.loads(req['body'])
model = body['params']['model']
# if model not known
if model not in MODELS:
logging.warning(f'Unknown model {model}')
continue
# if whitelist enabled and model not in it continue
if (len(self.model_whitelist) > 0 and
not model in self.model_whitelist):
continue
# if blacklist contains model skip
if model in self.model_blacklist:
continue
my_results = [res['id'] for res in self._snap['my_results']]
if rid not in my_results and rid in self._snap['requests']:
statuses = self._snap['requests'][rid]
if len(statuses) == 0:
binary, input_type = await self.conn.get_input_data(req['binary_data'])
hash_str = (
str(req['nonce'])
+
req['body']
+
req['binary_data']
)
logging.info(f'hashing: {hash_str}')
request_hash = sha256(hash_str.encode('utf-8')).hexdigest()
# TODO: validate request
# perform work
logging.info(f'working on {body}')
resp = await self.conn.begin_work(rid)
if 'code' in resp:
logging.info(f'probably being worked on already... skip.')
else:
try:
output_type = 'png'
if 'output_type' in body['params']:
output_type = body['params']['output_type']
output = None
output_hash = None
match self.backend:
case 'sync-on-thread':
self.mm._should_cancel = self.should_cancel_work
output_hash, output = await trio.to_thread.run_sync(
partial(
self.mm.compute_one,
rid,
body['method'], body['params'],
input_type=input_type,
binary=binary
)
)
case _:
raise DGPUComputeError(f'Unsupported backend {self.backend}')
self._last_generation_ts = datetime.now().isoformat()
self._last_benchmark = self._benchmark
self._benchmark = []
ipfs_hash = await self.conn.publish_on_ipfs(output, typ=output_type)
await self.conn.submit_work(rid, request_hash, output_hash, ipfs_hash)
except BaseException as e:
traceback.print_exc()
await self.conn.cancel_work(rid, str(e))
finally:
break
else:
logging.info(f'request {rid} already beign worked on, skip...')
if (await self.maybe_serve_one(req)):
break
await trio.sleep(1)

92
skynet/dgpu/network.py
View File

@ -8,15 +8,16 @@ import logging
from pathlib import Path
from functools import partial
import asks
import trio
import leap
import anyio
import httpx
from PIL import Image, UnidentifiedImageError
from leap.cleos import CLEOS
from leap.sugar import Checksum256, Name, asset_from_str
from skynet.constants import DEFAULT_IPFS_DOMAIN
from leap.protocol import Asset
from skynet.constants import DEFAULT_IPFS_DOMAIN, GPU_CONTRACT_ABI
from skynet.ipfs import AsyncIPFSHTTP, get_ipfs_file
from skynet.dgpu.errors import DGPUComputeError
@ -32,25 +33,25 @@ async def failable(fn: partial, ret_fail=None):
except (
OSError,
json.JSONDecodeError,
asks.errors.RequestTimeout,
asks.errors.BadHttpResponse,
anyio.BrokenResourceError
):
anyio.BrokenResourceError,
httpx.ReadError,
leap.errors.TransactionPushError
) as e:
return ret_fail
class SkynetGPUConnector:
def __init__(self, config: dict):
self.account = Name(config['account'])
self.account = config['account']
self.permission = config['permission']
self.key = config['key']
self.node_url = config['node_url']
self.hyperion_url = config['hyperion_url']
self.cleos = CLEOS(
None, None, self.node_url, remote=self.node_url)
self.cleos = CLEOS(endpoint=self.node_url)
self.cleos.load_abi('gpu.scd', GPU_CONTRACT_ABI)
self.ipfs_gateway_url = None
if 'ipfs_gateway_url' in config:
@ -151,11 +152,11 @@ class SkynetGPUConnector:
self.cleos.a_push_action,
'gpu.scd',
'workbegin',
{
list({
'worker': self.account,
'request_id': request_id,
'max_workers': 2
},
}.values()),
self.account, self.key,
permission=self.permission
)
@ -168,11 +169,11 @@ class SkynetGPUConnector:
self.cleos.a_push_action,
'gpu.scd',
'workcancel',
{
list({
'worker': self.account,
'request_id': request_id,
'reason': reason
},
}.values()),
self.account, self.key,
permission=self.permission
)
@ -191,10 +192,10 @@ class SkynetGPUConnector:
self.cleos.a_push_action,
'gpu.scd',
'withdraw',
{
list({
'user': self.account,
'quantity': asset_from_str(balance)
},
'quantity': Asset.from_str(balance)
}.values()),
self.account, self.key,
permission=self.permission
)
@ -226,13 +227,13 @@ class SkynetGPUConnector:
self.cleos.a_push_action,
'gpu.scd',
'submit',
{
list({
'worker': self.account,
'request_id': request_id,
'request_hash': Checksum256(request_hash),
'result_hash': Checksum256(result_hash),
'request_hash': request_hash,
'result_hash': result_hash,
'ipfs_hash': ipfs_hash
},
}.values()),
self.account, self.key,
permission=self.permission
)
@ -267,46 +268,15 @@ class SkynetGPUConnector:
return file_cid
async def get_input_data(self, ipfs_hash: str) -> tuple[bytes, str]:
input_type = 'none'
async def get_input_data(self, ipfs_hash: str) -> Image:
link = f'https://{self.ipfs_domain}/ipfs/{ipfs_hash}'
if ipfs_hash == '':
return b'', input_type
res = await get_ipfs_file(link, timeout=1)
logging.info(f'got response from {link}')
if not res or res.status_code != 200:
logging.warning(f'couldn\'t get ipfs binary data at {link}!')
results = {}
ipfs_link = f'https://{self.ipfs_domain}/ipfs/{ipfs_hash}'
ipfs_link_legacy = ipfs_link + '/image.png'
# attempt to decode as image
input_data = Image.open(io.BytesIO(res.raw))
async with trio.open_nursery() as n:
async def get_and_set_results(link: str):
res = await get_ipfs_file(link, timeout=1)
logging.info(f'got response from {link}')
if not res or res.status_code != 200:
logging.warning(f'couldn\'t get ipfs binary data at {link}!')
else:
try:
# attempt to decode as image
results[link] = Image.open(io.BytesIO(res.raw))
input_type = 'png'
n.cancel_scope.cancel()
except UnidentifiedImageError:
logging.warning(f'couldn\'t get ipfs binary data at {link}!')
n.start_soon(
get_and_set_results, ipfs_link)
n.start_soon(
get_and_set_results, ipfs_link_legacy)
input_data = None
if ipfs_link_legacy in results:
input_data = results[ipfs_link_legacy]
if ipfs_link in results:
input_data = results[ipfs_link]
if input_data == None:
raise DGPUComputeError('Couldn\'t gather input data from ipfs')
return input_data, input_type
return input_data

50
skynet/dgpu/pipes/flux.py 100644
View File

@ -0,0 +1,50 @@
#!/usr/bin/python
import torch
from diffusers import (
DiffusionPipeline,
FluxPipeline,
FluxTransformer2DModel
)
from transformers import T5EncoderModel, BitsAndBytesConfig
from huggingface_hub import hf_hub_download
__model = {
'name': 'black-forest-labs/FLUX.1-schnell'
}
def pipeline_for(
model: str,
mode: str,
mem_fraction: float = 1.0,
cache_dir: str | None = None
) -> DiffusionPipeline:
qonfig = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
params = {
'torch_dtype': torch.bfloat16,
'cache_dir': cache_dir,
'device_map': 'balanced',
'max_memory': {'cpu': '10GiB', 0: '11GiB'}
# 'max_memory': {0: '11GiB'}
}
text_encoder = T5EncoderModel.from_pretrained(
'black-forest-labs/FLUX.1-schnell',
subfolder="text_encoder_2",
torch_dtype=torch.bfloat16,
quantization_config=qonfig
)
params['text_encoder_2'] = text_encoder
pipe = FluxPipeline.from_pretrained(
model, **params)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
return pipe

56
skynet/dgpu/pipes/flux_inpaint.py 100644
View File

@ -0,0 +1,56 @@
#!/usr/bin/python
import torch
from diffusers import (
DiffusionPipeline,
FluxFillPipeline,
FluxTransformer2DModel
)
from transformers import T5EncoderModel, BitsAndBytesConfig
__model = {
'name': 'black-forest-labs/FLUX.1-Fill-dev'
}
def pipeline_for(
model: str,
mode: str,
mem_fraction: float = 1.0,
cache_dir: str | None = None
) -> DiffusionPipeline:
qonfig = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
params = {
'torch_dtype': torch.bfloat16,
'cache_dir': cache_dir,
'device_map': 'balanced',
'max_memory': {'cpu': '10GiB', 0: '11GiB'}
# 'max_memory': {0: '11GiB'}
}
text_encoder = T5EncoderModel.from_pretrained(
'sayakpaul/FLUX.1-Fill-dev-nf4',
subfolder="text_encoder_2",
torch_dtype=torch.bfloat16,
quantization_config=qonfig
)
params['text_encoder_2'] = text_encoder
transformer = FluxTransformer2DModel.from_pretrained(
'sayakpaul/FLUX.1-Fill-dev-nf4',
subfolder="transformer",
torch_dtype=torch.bfloat16,
quantization_config=qonfig
)
params['transformer'] = transformer
pipe = FluxFillPipeline.from_pretrained(
model, **params)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
return pipe

18
skynet/frontend/__init__.py
View File

@ -39,7 +39,7 @@ def validate_user_config_request(req: str):
case 'model' | 'algo':
attr = 'model'
val = params[2]
shorts = [model_info['short'] for model_info in MODELS.values()]
shorts = [model_info.short for model_info in MODELS.values()]
if val not in shorts:
raise ConfigUnknownAlgorithm(f'no model named {val}')
@ -112,20 +112,10 @@ def validate_user_config_request(req: str):
def perform_auto_conf(config: dict) -> dict:
model = config['model']
prefered_size_w = 512
prefered_size_h = 512
if 'xl' in model:
prefered_size_w = 1024
prefered_size_h = 1024
else:
prefered_size_w = 512
prefered_size_h = 512
model = MODELS[config['model']]
config['step'] = random.randint(20, 35)
config['width'] = prefered_size_w
config['height'] = prefered_size_h
config['width'] = model.size.w
config['height'] = model.size.h
return config

76
skynet/frontend/telegram/__init__.py
View File

@ -14,7 +14,7 @@ from contextlib import AsyncExitStack
from contextlib import asynccontextmanager as acm
from leap.cleos import CLEOS
from leap.sugar import Name, asset_from_str, collect_stdout
from leap.protocol import Name, Asset
from leap.hyperion import HyperionAPI
from telebot.types import InputMediaPhoto
@ -43,7 +43,6 @@ class SkynetTelegramFrontend:
db_user: str,
db_pass: str,
ipfs_node: str,
remote_ipfs_node: str | None,
key: str,
explorer_domain: str,
ipfs_domain: str
@ -56,22 +55,19 @@ class SkynetTelegramFrontend:
self.db_host = db_host
self.db_user = db_user
self.db_pass = db_pass
self.remote_ipfs_node = remote_ipfs_node
self.key = key
self.explorer_domain = explorer_domain
self.ipfs_domain = ipfs_domain
self.bot = AsyncTeleBot(token, exception_handler=SKYExceptionHandler)
self.cleos = CLEOS(None, None, url=node_url, remote=node_url)
self.cleos = CLEOS(endpoint=node_url)
self.cleos.load_abi('gpu.scd', GPU_CONTRACT_ABI)
self.hyperion = HyperionAPI(hyperion_url)
self.ipfs_node = AsyncIPFSHTTP(ipfs_node)
self._async_exit_stack = AsyncExitStack()
async def start(self):
if self.remote_ipfs_node:
await self.ipfs_node.connect(self.remote_ipfs_node)
self.db_call = await self._async_exit_stack.enter_async_context(
open_database_connection(
self.db_user, self.db_pass, self.db_host))
@ -116,7 +112,7 @@ class SkynetTelegramFrontend:
method: str,
params: dict,
file_id: str | None = None,
binary_data: str = ''
inputs: list[str] = []
) -> bool:
if params['seed'] == None:
params['seed'] = random.randint(0, 0xFFFFFFFF)
@ -145,13 +141,13 @@ class SkynetTelegramFrontend:
res = await self.cleos.a_push_action(
'gpu.scd',
'enqueue',
{
list({
'user': Name(self.account),
'request_body': body,
'binary_data': binary_data,
'reward': asset_from_str(reward),
'binary_data': ','.join(inputs),
'reward': Asset.from_str(reward),
'min_verification': 1
},
}.values()),
self.account, self.key, permission=self.permission
)
@ -176,12 +172,12 @@ class SkynetTelegramFrontend:
parse_mode='HTML'
)
out = collect_stdout(res)
out = res['processed']['action_traces'][0]['console']
request_id, nonce = out.split(':')
request_hash = sha256(
(nonce + body + binary_data).encode('utf-8')).hexdigest().upper()
(nonce + body + ','.join(inputs)).encode('utf-8')).hexdigest().upper()
request_id = int(request_id)
@ -189,7 +185,7 @@ class SkynetTelegramFrontend:
tx_hash = None
ipfs_hash = None
for i in range(60):
for i in range(60 * 3):
try:
submits = await self.hyperion.aget_actions(
account=self.account,
@ -241,46 +237,28 @@ class SkynetTelegramFrontend:
user, params, tx_hash, worker, reward, self.explorer_domain)
# attempt to get the image and send it
results = {}
ipfs_link = f'https://{self.ipfs_domain}/ipfs/{ipfs_hash}'
ipfs_link_legacy = ipfs_link + '/image.png'
async def get_and_set_results(link: str):
res = await get_ipfs_file(link)
logging.info(f'got response from {link}')
if not res or res.status_code != 200:
logging.warning(f'couldn\'t get ipfs binary data at {link}!')
res = await get_ipfs_file(ipfs_link)
logging.info(f'got response from {ipfs_link}')
if not res or res.status_code != 200:
logging.warning(f'couldn\'t get ipfs binary data at {ipfs_link}!')
else:
try:
with Image.open(io.BytesIO(res.raw)) as image:
w, h = image.size
else:
try:
with Image.open(io.BytesIO(res.raw)) as image:
w, h = image.size
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
logging.warning(f'result is of size {image.size}')
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
logging.warning(f'result is of size {image.size}')
image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
tmp_buf = io.BytesIO()
image.save(tmp_buf, format='PNG')
png_img = tmp_buf.getvalue()
tmp_buf = io.BytesIO()
image.save(tmp_buf, format='PNG')
png_img = tmp_buf.getvalue()
results[link] = png_img
except UnidentifiedImageError:
logging.warning(f'couldn\'t get ipfs binary data at {link}!')
tasks = [
get_and_set_results(ipfs_link),
get_and_set_results(ipfs_link_legacy)
]
await asyncio.gather(*tasks)
png_img = None
if ipfs_link_legacy in results:
png_img = results[ipfs_link_legacy]
if ipfs_link in results:
png_img = results[ipfs_link]
except UnidentifiedImageError:
logging.warning(f'couldn\'t get ipfs binary data at {ipfs_link}!')
if not png_img:
await self.update_status_message(

4
skynet/frontend/telegram/handlers.py
View File

@ -254,7 +254,7 @@ def create_handler_context(frontend: 'SkynetTelegramFrontend'):
success = await work_request(
user, status_msg, 'img2img', params,
file_id=file_id,
binary_data=ipfs_hash
inputs=ipfs_hash
)
if success:
@ -320,7 +320,7 @@ def create_handler_context(frontend: 'SkynetTelegramFrontend'):
success = await work_request(
user, status_msg, 'redo', params,
file_id=file_id,
binary_data=binary
inputs=binary
)
if success:

2
skynet/frontend/telegram/utils.py
View File

@ -72,7 +72,7 @@ def generate_reply_caption(
):
explorer_link = hlink(
'SKYNET Transaction Explorer',
f'https://explorer.{explorer_domain}/v2/explore/transaction/{tx_hash}'
f'https://{explorer_domain}/v2/explore/transaction/{tx_hash}'
)
meta_info = prepare_metainfo_caption(tguser, worker, reward, params)

32
skynet/ipfs/__init__.py
View File

@ -3,10 +3,10 @@
import logging
from pathlib import Path
import asks
import httpx
class IPFSClientException(BaseException):
class IPFSClientException(Exception):
...
@ -16,19 +16,20 @@ class AsyncIPFSHTTP:
self.endpoint = endpoint
async def _post(self, sub_url: str, *args, **kwargs):
resp = await asks.post(
self.endpoint + sub_url,
*args, **kwargs
)
async with httpx.AsyncClient() as client:
resp = await client.post(
self.endpoint + sub_url,
*args, **kwargs
)
if resp.status_code != 200:
raise IPFSClientException(resp.text)
return resp.json()
#!/usr/bin/python
async def add(self, file_path: Path, **kwargs):
files = {
'file': file_path
'file': (file_path.name, file_path.open('rb'))
}
return await self._post(
'/api/v0/add',
@ -55,18 +56,19 @@ class AsyncIPFSHTTP:
))['Peers']
async def get_ipfs_file(ipfs_link: str, timeout: int = 60):
async def get_ipfs_file(ipfs_link: str, timeout: int = 60 * 5):
logging.info(f'attempting to get image at {ipfs_link}')
resp = None
for i in range(timeout):
for _ in range(timeout):
try:
resp = await asks.get(ipfs_link, timeout=3)
async with httpx.AsyncClient() as client:
resp = await client.get(ipfs_link, timeout=3)
except asks.errors.RequestTimeout:
logging.warning('timeout...')
except httpx.RequestError as e:
logging.warning(f'Request error: {e}')
except asks.errors.BadHttpResponse as e:
logging.error(f'ifps gateway exception: \n{e}')
if resp is not None:
break
if resp:
logging.info(f'status_code: {resp.status_code}')

6
skynet/ipfs/pinner.py
View File

@ -55,9 +55,9 @@ class SkynetPinner:
cids = []
for action in enqueues['actions']:
cid = action['act']['data']['binary_data']
if cid and not self.is_pinned(cid):
cids.append(cid)
for cid in action['act']['data']['binary_data'].split(','):
if cid and not self.is_pinned(cid):
cids.append(cid)
return cids

164
skynet/utils.py
View File

@ -6,26 +6,42 @@ import sys
import time
import random
import logging
import importlib
from typing import Optional
from pathlib import Path
import asks
import trio
import torch
import numpy as np
from PIL import Image
from basicsr.archs.rrdbnet_arch import RRDBNet
from diffusers import (
DiffusionPipeline,
EulerAncestralDiscreteScheduler
AutoPipelineForText2Image,
AutoPipelineForImage2Image,
AutoPipelineForInpainting,
EulerAncestralDiscreteScheduler,
)
from realesrgan import RealESRGANer
from huggingface_hub import login
import trio
from .constants import MODELS
# Hack to fix a changed import in torchvision 0.17+, which otherwise breaks
# basicsr; see https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13985
try:
import torchvision.transforms.functional_tensor # noqa: F401
except ImportError:
try:
import torchvision.transforms.functional as functional
sys.modules["torchvision.transforms.functional_tensor"] = functional
except ImportError:
pass # shrug...
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
def time_ms():
return int(time.time() * 1000)
@ -58,14 +74,18 @@ def crop_image(image: Image, max_w: int, max_h: int) -> Image:
return image.convert('RGB')
def convert_from_bytes_and_crop(raw: bytes, max_w: int, max_h: int) -> Image:
return crop_image(convert_from_bytes_to_img(raw), max_w, max_h)
def pipeline_for(
model: str,
mode: str,
mem_fraction: float = 1.0,
image: bool = False,
cache_dir: str | None = None
) -> DiffusionPipeline:
logging.info(f'pipeline_for {model} {mode}')
assert torch.cuda.is_available()
torch.cuda.empty_cache()
torch.backends.cuda.matmul.allow_tf32 = True
@ -79,21 +99,35 @@ def pipeline_for(
torch.use_deterministic_algorithms(True)
model_info = MODELS[model]
shortname = model_info.short
# disable for compat with "diffuse" method
# assert mode in model_info.tags
# default to checking if custom pipeline exist and return that if not, attempt generic
try:
normalized_shortname = shortname.replace('-', '_')
custom_pipeline = importlib.import_module(f'skynet.dgpu.pipes.{normalized_shortname}')
assert custom_pipeline.__model['name'] == model
return custom_pipeline.pipeline_for(model, mode, mem_fraction=mem_fraction, cache_dir=cache_dir)
except ImportError:
...
req_mem = model_info.mem
req_mem = model_info['mem']
mem_gb = torch.cuda.mem_get_info()[1] / (10**9)
mem_gb *= mem_fraction
over_mem = mem_gb < req_mem
if over_mem:
logging.warn(f'model requires {req_mem} but card has {mem_gb}, model will run slower..')
shortname = model_info['short']
params = {
'safety_checker': None,
'torch_dtype': torch.float16,
'cache_dir': cache_dir,
'variant': 'fp16'
'variant': 'fp16',
}
match shortname:
@ -102,26 +136,37 @@ def pipeline_for(
torch.cuda.set_per_process_memory_fraction(mem_fraction)
pipe = DiffusionPipeline.from_pretrained(
pipe_class = DiffusionPipeline
match mode:
case 'inpaint':
pipe_class = AutoPipelineForInpainting
case 'img2img':
pipe_class = AutoPipelineForImage2Image
case 'txt2img':
pipe_class = AutoPipelineForText2Image
pipe = pipe_class.from_pretrained(
model, **params)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_xformers_memory_efficient_attention()
if over_mem:
if not image:
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()
if mode == 'txt2img':
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.enable_model_cpu_offload()
else:
if sys.version_info[1] < 11:
# torch.compile only supported on python < 3.11
pipe.unet = torch.compile(
pipe.unet, mode='reduce-overhead', fullgraph=True)
# if sys.version_info[1] < 11:
# # torch.compile only supported on python < 3.11
# pipe.unet = torch.compile(
# pipe.unet, mode='reduce-overhead', fullgraph=True)
pipe = pipe.to('cuda')
@ -130,7 +175,7 @@ def pipeline_for(
def txt2img(
hf_token: str,
model: str = 'prompthero/openjourney',
model: str = list(MODELS.keys())[-1],
prompt: str = 'a red old tractor in a sunny wheat field',
output: str = 'output.png',
width: int = 512, height: int = 512,
@ -139,7 +184,7 @@ def txt2img(
seed: Optional[int] = None
):
login(token=hf_token)
pipe = pipeline_for(model)
pipe = pipeline_for(model, 'txt2img')
seed = seed if seed else random.randint(0, 2 ** 64)
prompt = prompt
@ -156,7 +201,7 @@ def txt2img(
def img2img(
hf_token: str,
model: str = 'prompthero/openjourney',
model: str = list(MODELS.keys())[-2],
prompt: str = 'a red old tractor in a sunny wheat field',
img_path: str = 'input.png',
output: str = 'output.png',
@ -166,10 +211,12 @@ def img2img(
seed: Optional[int] = None
):
login(token=hf_token)
pipe = pipeline_for(model, image=True)
pipe = pipeline_for(model, 'img2img')
model_info = MODELS[model]
with open(img_path, 'rb') as img_file:
input_img = convert_from_bytes_and_crop(img_file.read(), 512, 512)
input_img = convert_from_bytes_and_crop(img_file.read(), model_info.size.w, model_info.size.h)
seed = seed if seed else random.randint(0, 2 ** 64)
prompt = prompt
@ -184,7 +231,48 @@ def img2img(
image.save(output)
def init_upscaler(model_path: str = 'weights/RealESRGAN_x4plus.pth'):
def inpaint(
hf_token: str,
model: str = list(MODELS.keys())[-3],
prompt: str = 'a red old tractor in a sunny wheat field',
img_path: str = 'input.png',
mask_path: str = 'mask.png',
output: str = 'output.png',
strength: float = 1.0,
guidance: float = 10,
steps: int = 28,
seed: Optional[int] = None
):
login(token=hf_token)
pipe = pipeline_for(model, 'inpaint')
model_info = MODELS[model]
with open(img_path, 'rb') as img_file:
input_img = convert_from_bytes_and_crop(img_file.read(), model_info.size.w, model_info.size.h)
with open(mask_path, 'rb') as mask_file:
mask_img = convert_from_bytes_and_crop(mask_file.read(), model_info.size.w, model_info.size.h)
var_params = {}
if 'flux' not in model.lower():
var_params['strength'] = strength
seed = seed if seed else random.randint(0, 2 ** 64)
prompt = prompt
image = pipe(
prompt,
image=input_img,
mask_image=mask_img,
guidance_scale=guidance, num_inference_steps=steps,
generator=torch.Generator("cuda").manual_seed(seed),
**var_params
).images[0]
image.save(output)
def init_upscaler(model_path: str = 'hf_home/RealESRGAN_x4plus.pth'):
return RealESRGANer(
scale=4,
model_path=model_path,
@ -203,7 +291,7 @@ def init_upscaler(model_path: str = 'weights/RealESRGAN_x4plus.pth'):
def upscale(
img_path: str = 'input.png',
output: str = 'output.png',
model_path: str = 'weights/RealESRGAN_x4plus.pth'
model_path: str = 'hf_home/RealESRGAN_x4plus.pth'
):
input_img = Image.open(img_path).convert('RGB')
@ -214,25 +302,3 @@ def upscale(
image = convert_from_cv2_to_image(up_img)
image.save(output)
async def download_upscaler():
print('downloading upscaler...')
weights_path = Path('weights')
weights_path.mkdir(exist_ok=True)
upscaler_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
save_path = weights_path / 'RealESRGAN_x4plus.pth'
response = await asks.get(upscaler_url)
with open(save_path, 'wb') as f:
f.write(response.content)
print('done')
def download_all_models(hf_token: str, hf_home: str):
assert torch.cuda.is_available()
trio.run(download_upscaler)
login(token=hf_token)
for model in MODELS:
print(f'DOWNLOADING {model.upper()}')
pipeline_for(model, cache_dir=hf_home)

20
tests/conftest.py
View File

@ -2,7 +2,7 @@
import pytest
from skynet.db import open_new_database
from skynet.config import *
from skynet.ipfs import AsyncIPFSHTTP
from skynet.ipfs.docker import open_ipfs_node
from skynet.nodeos import open_nodeos
@ -15,6 +15,7 @@ def ipfs_client():
@pytest.fixture(scope='session')
def postgres_db():
from skynet.db import open_new_database
with open_new_database() as db_params:
yield db_params
@ -22,3 +23,20 @@ def postgres_db():
def cleos():
with open_nodeos() as cli:
yield cli
@pytest.fixture(scope='session')
def dgpu():
from skynet.dgpu.network import SkynetGPUConnector
from skynet.dgpu.compute import SkynetMM
from skynet.dgpu.daemon import SkynetDGPUDaemon
config = load_skynet_toml(file_path='skynet.toml')
hf_token = load_key(config, 'skynet.dgpu.hf_token')
hf_home = load_key(config, 'skynet.dgpu.hf_home')
set_hf_vars(hf_token, hf_home)
config = config['skynet']['dgpu']
conn = SkynetGPUConnector(config)
mm = SkynetMM(config)
daemon = SkynetDGPUDaemon(mm, conn, config)
yield conn, mm, daemon

112
tests/test_reqs.py 100644
View File

@ -0,0 +1,112 @@
import json
from skynet.dgpu.compute import SkynetMM
from skynet.constants import *
from skynet.config import *
async def test_diffuse(dgpu):
conn, mm, daemon = dgpu
await conn.cancel_work(0, 'testing')
daemon._snap['requests'][0] = {}
req = {
'id': 0,
'nonce': 0,
'body': json.dumps({
"method": "diffuse",
"params": {
"prompt": "Kronos God Realistic 4k",
"model": list(MODELS.keys())[-1],
"step": 21,
"width": 1024,
"height": 1024,
"seed": 168402949,
"guidance": "7.5"
}
}),
'binary_data': '',
}
await daemon.maybe_serve_one(req)
async def test_txt2img(dgpu):
conn, mm, daemon = dgpu
await conn.cancel_work(0, 'testing')
daemon._snap['requests'][0] = {}
req = {
'id': 0,
'nonce': 0,
'body': json.dumps({
"method": "txt2img",
"params": {
"prompt": "Kronos God Realistic 4k",
"model": list(MODELS.keys())[-1],
"step": 21,
"width": 1024,
"height": 1024,
"seed": 168402949,
"guidance": "7.5"
}
}),
'binary_data': '',
}
await daemon.maybe_serve_one(req)
async def test_img2img(dgpu):
conn, mm, daemon = dgpu
await conn.cancel_work(0, 'testing')
daemon._snap['requests'][0] = {}
req = {
'id': 0,
'nonce': 0,
'body': json.dumps({
"method": "img2img",
"params": {
"prompt": "a hindu cat god feline god on a house roof",
"model": list(MODELS.keys())[-2],
"step": 21,
"width": 1024,
"height": 1024,
"seed": 168402949,
"guidance": "7.5",
"strength": "0.5"
}
}),
'binary_data': 'QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi',
}
await daemon.maybe_serve_one(req)
async def test_inpaint(dgpu):
conn, mm, daemon = dgpu
await conn.cancel_work(0, 'testing')
daemon._snap['requests'][0] = {}
req = {
'id': 0,
'nonce': 0,
'body': json.dumps({
"method": "inpaint",
"params": {
"prompt": "a black panther on a sunny roof",
"model": list(MODELS.keys())[-3],
"step": 21,
"width": 1024,
"height": 1024,
"seed": 168402949,
"guidance": "7.5",
"strength": "0.5"
}
}),
'binary_data':
'QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi,' +
'Qmccx1aXNmq5mZDS3YviUhgGHXWhQeHvca3AgA7MDjj2hR'
}
await daemon.maybe_serve_one(req)