diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b031f5..607a122 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,6 +26,4 @@ jobs: - name: Run tests run: | - uv run \ - pytest \ - tests/test_chain.py + uv run pytest diff --git a/skynet/dgpu/compute.py b/skynet/dgpu/compute.py index 2856a60..f5892f3 100755 --- a/skynet/dgpu/compute.py +++ b/skynet/dgpu/compute.py @@ -32,25 +32,25 @@ def prepare_params_for_diffuse( match mode: case ModelMode.INPAINT: image = crop_image( - inputs[0], params['width'], params['height']) + inputs[0], params.width, params.height) mask = crop_image( - inputs[1], params['width'], params['height']) + inputs[1], params.width, params.height) _params['image'] = image _params['mask_image'] = mask - if 'flux' in params['model'].lower(): + if 'flux' in params.model.lower(): _params['max_sequence_length'] = 512 else: - _params['strength'] = float(params['strength']) + _params['strength'] = params.strength case ModelMode.IMG2IMG: image = crop_image( - inputs[0], params['width'], params['height']) + inputs[0], params.width, params.height) _params['image'] = image - _params['strength'] = float(params['strength']) + _params['strength'] = params.strength case ModelMode.TXT2IMG | ModelMode.DIFFUSE: ... @@ -63,7 +63,6 @@ def prepare_params_for_diffuse( params.guidance, params.step, torch.manual_seed(int(params.seed)), - params.upscaler, _params ) @@ -88,12 +87,8 @@ def maybe_load_model(name: str, mode: ModelMode): _model_name = _model_mode = '' # load model - if mode == ModelMode.UPSCALE: - _model = init_upscaler() - - else: - _model = pipeline_for( - name, mode, cache_dir=config.hf_home) + _model = pipeline_for( + name, mode, cache_dir=config.hf_home) _model_name = name _model_mode = mode @@ -154,7 +149,7 @@ def compute_one( ): arguments = prepare_params_for_diffuse( params, method, inputs) - prompt, guidance, step, seed, upscaler, extra_params = arguments + prompt, guidance, step, seed, extra_params = arguments if 'flux' in name.lower(): extra_params['callback_on_step_end'] = inference_step_wakeup @@ -174,13 +169,6 @@ def compute_one( output_binary = b'' match output_type: case 'png': - if upscaler == 'x4': - input_img = output.convert('RGB') - up_img, _ = init_upscaler().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) case _: diff --git a/skynet/dgpu/pipes/__init__.py b/skynet/dgpu/pipes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/skynet/dgpu/pipes/tester.py b/skynet/dgpu/pipes/tester.py index 7978f22..58c6181 100644 --- a/skynet/dgpu/pipes/tester.py +++ b/skynet/dgpu/pipes/tester.py @@ -6,7 +6,7 @@ import msgspec __model = { - 'name': 'skygpu/txt2img-mocker' + 'name': 'skygpu/mocker' } class MockPipelineResult(msgspec.Struct): diff --git a/skynet/dgpu/pipes/tester_upscale.py b/skynet/dgpu/pipes/tester_upscale.py new file mode 100644 index 0000000..934802a --- /dev/null +++ b/skynet/dgpu/pipes/tester_upscale.py @@ -0,0 +1,34 @@ +from PIL import Image + +import msgspec + +from skynet.dgpu.utils import convert_from_image_to_cv2 + + +__model = { + 'name': 'skygpu/mocker-upscale' +} + +class MockPipelineResult(msgspec.Struct): + images: list[Image] + +class MockUpscalePipeline: + + def enhance( + self, + img, + outscale: int + ): + img = Image.new('RGB', (outscale, outscale), color='green') + output = convert_from_image_to_cv2(img) + + return (output, None) + + +def pipeline_for( + model: str, + mode: str, + mem_fraction: float = 1.0, + cache_dir: str | None = None +): + return MockUpscalePipeline() diff --git a/skynet/dgpu/utils.py b/skynet/dgpu/utils.py index 6533d26..6a17c93 100755 --- a/skynet/dgpu/utils.py +++ b/skynet/dgpu/utils.py @@ -124,6 +124,9 @@ def pipeline_for( except ImportError: logging.info(f'didn\'t find a custom pipeline file for {shortname}') + # for now, upscaler special case... + if mode == 'upscale': + return init_upscaler() req_mem = model_info.mem diff --git a/tests/conftest.py b/tests/conftest.py index 1540f9d..7cb2fc0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import pytest from skynet.ipfs import AsyncIPFSHTTP +from skynet._testing import override_dgpu_config @pytest.fixture(scope='session') @@ -44,13 +45,35 @@ def skynet_cleos(cleos_bs): @pytest.fixture def inject_mockers(): from skynet.constants import MODELS - from skynet.types import ModelDesc + from skynet.types import ModelDesc, ModelMode - MODELS['skygpu/txt2img-mocker'] = ModelDesc( + MODELS['skygpu/mocker'] = ModelDesc( short='tester', mem=0.01, attrs={}, - tags=['txt2img'] + tags=[ + ModelMode.TXT2IMG, + ModelMode.IMG2IMG, + ModelMode.INPAINT + ] + ) + + MODELS['skygpu/mocker-upscale'] = ModelDesc( + short='tester-upscale', + mem=0.01, + attrs={}, + tags=[ + ModelMode.UPSCALE + ] + ) + + override_dgpu_config( + account='testworker1', + permission='active', + key='', + node_url='', + ipfs_url='http://127.0.0.1:5001', + hf_token='' ) yield diff --git a/tests/test_chain.py b/tests/test_chain.py index 124ee8f..97adfe6 100644 --- a/tests/test_chain.py +++ b/tests/test_chain.py @@ -19,7 +19,7 @@ async def test_full_flow(inject_mockers, skynet_cleos, ipfs_node): method='txt2img', params=BodyV0Params( prompt='cyberpunk hacker travis bickle dystopic alley graffiti', - model='skygpu/txt2img-mocker', + model='skygpu/mocker', step=4, seed=0, guidance=10.0 diff --git a/tests/test_ipfs_client.py b/tests/test_ipfs_client.py index a400cdf..c857392 100644 --- a/tests/test_ipfs_client.py +++ b/tests/test_ipfs_client.py @@ -1,14 +1,16 @@ from pathlib import Path -async def test_connection(ipfs_client): - await ipfs_client.connect( - '/ip4/169.197.140.154/tcp/4001/p2p/12D3KooWKWogLFNEcNNMKnzU7Snrnuj84RZdMBg3sLiQSQc51oEv') - peers = await ipfs_client.peers() - assert '12D3KooWKWogLFNEcNNMKnzU7Snrnuj84RZdMBg3sLiQSQc51oEv' in [p['Peer'] for p in peers] +# async def test_connection(ipfs_node): +# _, ipfs_client = ipfs_node +# await ipfs_client.connect( +# '/ip4/169.197.140.154/tcp/4001/p2p/12D3KooWKWogLFNEcNNMKnzU7Snrnuj84RZdMBg3sLiQSQc51oEv') +# peers = await ipfs_client.peers() +# assert '12D3KooWKWogLFNEcNNMKnzU7Snrnuj84RZdMBg3sLiQSQc51oEv' in [p['Peer'] for p in peers] -async def test_add_and_pin_file(ipfs_client): +async def test_add_and_pin_file(ipfs_node): + _, ipfs_client = ipfs_node test_file = Path('hello_world.txt') with open(test_file, 'w+') as file: file.write('Hello Skynet!') diff --git a/tests/test_reqs.py b/tests/test_reqs.py index 274c691..a5b9e30 100644 --- a/tests/test_reqs.py +++ b/tests/test_reqs.py @@ -1,24 +1,23 @@ import pytest +from PIL import Image from skynet.types import ModelMode, BodyV0Params from skynet.dgpu.compute import maybe_load_model, compute_one -from skynet._testing import override_dgpu_config - -@pytest.mark.parametrize("mode", [ - (ModelMode.DIFFUSE), (ModelMode.TXT2IMG) +@pytest.mark.parametrize('mode,model', [ + (ModelMode.DIFFUSE, 'skygpu/mocker'), + (ModelMode.TXT2IMG, 'skygpu/mocker'), + (ModelMode.IMG2IMG, 'skygpu/mocker'), + (ModelMode.INPAINT, 'skygpu/mocker'), + (ModelMode.UPSCALE, 'skygpu/mocker-upscale'), ]) -async def test_pipeline_mocker(inject_mockers, mode): - override_dgpu_config( - account='testworker1', - permission='active', - key='', - node_url='', - ipfs_url='http://127.0.0.1:5001', - hf_token='' - ) - model = 'skygpu/txt2img-mocker' +async def test_pipeline_mocker(inject_mockers, mode, model): + # always insert at least two inputs to make all modes pass + inputs = [ + Image.new('RGB', (1, 1), color='green') + for i in range(2) + ] params = BodyV0Params( prompt="Kronos God Realistic 4k", model=model, @@ -30,21 +29,21 @@ async def test_pipeline_mocker(inject_mockers, mode): ) with maybe_load_model(model, mode) as model: - compute_one(model, 0, mode, params) + compute_one(model, 0, mode, params, inputs) - -async def test_pipeline(): - model = 'stabilityai/stable-diffusion-xl-base-1.0' - mode = 'txt2img' - params = BodyV0Params( - prompt="Kronos God Realistic 4k", - model=model, - step=21, - width=1024, - height=1024, - seed=168402949, - guidance="7.5" - ) - - with maybe_load_model(model, mode) as model: - compute_one(model, 0, mode, params) +# disable for now (cuda) +# async def test_pipeline(): +# model = 'stabilityai/stable-diffusion-xl-base-1.0' +# mode = 'txt2img' +# params = BodyV0Params( +# prompt="Kronos God Realistic 4k", +# model=model, +# step=21, +# width=1024, +# height=1024, +# seed=168402949, +# guidance="7.5" +# ) +# +# with maybe_load_model(model, mode) as model: +# compute_one(model, 0, mode, params)