From 22e40b766f72ff63d83d3559be39135c9ebb2fc7 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 9 Jan 2025 21:25:00 -0300 Subject: [PATCH] Add dgpu fixture --- skynet/dgpu/compute.py | 1 + skynet/dgpu/daemon.py | 1 + skynet/dgpu/network.py | 2 +- skynet/utils.py | 4 +-- tests/conftest.py | 20 ++++++++++++- tests/test_reqs.py | 67 +++++++++++++++++------------------------- 6 files changed, 51 insertions(+), 44 deletions(-) diff --git a/skynet/dgpu/compute.py b/skynet/dgpu/compute.py index a0cbe62..049f00b 100644 --- a/skynet/dgpu/compute.py +++ b/skynet/dgpu/compute.py @@ -34,6 +34,7 @@ def prepare_params_for_diffuse( inputs[1], params['width'], params['height']) _params['image'] = image + _params['mask_image'] = mask _params['strength'] = float(params['strength']) case 'img2img': diff --git a/skynet/dgpu/daemon.py b/skynet/dgpu/daemon.py index 34b574c..005c2a1 100644 --- a/skynet/dgpu/daemon.py +++ b/skynet/dgpu/daemon.py @@ -146,6 +146,7 @@ class SkynetDGPUDaemon: inputs = [ await self.conn.get_input_data(_input) for _input in req['binary_data'].split(',') + if _input ] hash_str = ( diff --git a/skynet/dgpu/network.py b/skynet/dgpu/network.py index 46f1363..1e8315d 100644 --- a/skynet/dgpu/network.py +++ b/skynet/dgpu/network.py @@ -268,7 +268,7 @@ class SkynetGPUConnector: return file_cid async def get_input_data(self, ipfs_hash: str) -> Image: - ipfs_link = f'https://{self.ipfs_domain}/ipfs/{ipfs_hash}' + link = f'https://{self.ipfs_domain}/ipfs/{ipfs_hash}' res = await get_ipfs_file(link, timeout=1) logging.info(f'got response from {link}') diff --git a/skynet/utils.py b/skynet/utils.py index 9ea2cf0..0b43451 100755 --- a/skynet/utils.py +++ b/skynet/utils.py @@ -65,8 +65,8 @@ def convert_from_bytes_and_crop(raw: bytes, max_w: int, max_h: int) -> Image: def pipeline_for( model: str, + mode: str, mem_fraction: float = 1.0, - mode: str = [], cache_dir: str | None = None ) -> DiffusionPipeline: @@ -112,7 +112,7 @@ def pipeline_for( else: pipe_class = DiffusionPipeline - pipe = AutoPipelineForInpainting.from_pretrained( + pipe = pipe_class.from_pretrained( model, **params) pipe.scheduler = EulerAncestralDiscreteScheduler.from_config( diff --git a/tests/conftest.py b/tests/conftest.py index 03ed434..0ea4821 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_reqs.py b/tests/test_reqs.py index cf97cee..fbd5f73 100644 --- a/tests/test_reqs.py +++ b/tests/test_reqs.py @@ -1,9 +1,17 @@ +import json +from skynet.dgpu.compute import SkynetMM +from skynet.constants import * from skynet.config import * -async def test_txt2img(): +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": { @@ -16,25 +24,20 @@ async def test_txt2img(): "guidance": "7.5" } }), - 'inputs': [], + 'binary_data': '', } - config = load_skynet_toml(file_path=config_path) - 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) - - assert 'skynet' in config - assert 'dgpu' in config['skynet'] - - mm = SkynetMM(config['skynet']['dgpu']) - - mm.maybe_serve_one(req) + await daemon.maybe_serve_one(req) -async def test_img2img(): +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": { @@ -48,24 +51,19 @@ async def test_img2img(): "strength": "0.5" } }), - 'inputs': ['QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi'], + 'binary_data': 'QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi', } - config = load_skynet_toml(file_path=config_path) - 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) + await daemon.maybe_serve_one(req) - assert 'skynet' in config - assert 'dgpu' in config['skynet'] +async def test_inpaint(dgpu): + conn, mm, daemon = dgpu + await conn.cancel_work(0, 'testing') - mm = SkynetMM(config['skynet']['dgpu']) - - mm.maybe_serve_one(req) - -async def test_inpaint(): + daemon._snap['requests'][0] = {} req = { 'id': 0, + 'nonce': 0, 'body': json.dumps({ "method": "inpaint", "params": { @@ -79,20 +77,9 @@ async def test_inpaint(): "strength": "0.5" } }), - 'inputs': [ - 'QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi', + 'binary_data': + 'QmZcGdXXVQfpco1G3tr2CGFBtv8xVsCwcwuq9gnJBWDymi,' + 'Qmccx1aXNmq5mZDS3YviUhgGHXWhQeHvca3AgA7MDjj2hR' - ], } - config = load_skynet_toml(file_path=config_path) - 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) - - assert 'skynet' in config - assert 'dgpu' in config['skynet'] - - mm = SkynetMM(config['skynet']['dgpu']) - - mm.maybe_serve_one(req) + await daemon.maybe_serve_one(req)