From 8427165a76335124909f5961a856e9572351dd55 Mon Sep 17 00:00:00 2001
From: Guillermo Rodriguez <guillermo@telos.net>
Date: Tue, 17 Jan 2023 19:01:31 -0300
Subject: [PATCH] Add image resize before img2img calc, add strength parameter
 to local img2img tools

---
 skynet/cli.py   | 4 +++-
 skynet/dgpu.py  | 7 +++++++
 skynet/utils.py | 2 ++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/skynet/cli.py b/skynet/cli.py
index de7a37d..2573106 100644
--- a/skynet/cli.py
+++ b/skynet/cli.py
@@ -47,10 +47,11 @@ def txt2img(*args, **kwargs):
     '--prompt', '-p', default='a red old tractor in a sunny wheat field')
 @click.option('--input', '-i', default='input.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 img2img(model, prompt, input, output, guidance, steps, seed):
+def img2img(model, prompt, input, output, strength, guidance, steps, seed):
     assert 'HF_TOKEN' in os.environ
     utils.img2img(
         os.environ['HF_TOKEN'],
@@ -58,6 +59,7 @@ def img2img(model, prompt, input, output, guidance, steps, seed):
         prompt=prompt,
         img_path=input,
         output=output,
+        strength=strength,
         guidance=guidance,
         steps=steps,
         seed=seed
diff --git a/skynet/dgpu.py b/skynet/dgpu.py
index 2a8118e..1520c06 100644
--- a/skynet/dgpu.py
+++ b/skynet/dgpu.py
@@ -279,6 +279,13 @@ async def open_dgpu_node(
                     raw_img = zlib.decompress(img_raw)
                     logging.info(raw_img[:10])
                     img = Image.open(io.BytesIO(raw_img))
+                    w, h = img.size
+                    logging.info(f'user sent img of size {img.size}')
+
+                    if w > 512 or h > 512:
+                        img.thumbnail((512, 512))
+                        logging.info(f'resized it to {img.size}')
+
 
                 req = DGPUBusMessage()
                 req.ParseFromString(msg)
diff --git a/skynet/utils.py b/skynet/utils.py
index 64a4583..ba1ce2d 100644
--- a/skynet/utils.py
+++ b/skynet/utils.py
@@ -101,6 +101,7 @@ def img2img(
     prompt: str = 'a red old tractor in a sunny wheat field',
     img_path: str = 'input.png',
     output: str = 'output.png',
+    strength: float = 1.0,
     guidance: float = 10,
     steps: int = 28,
     seed: Optional[int] = None
@@ -121,6 +122,7 @@ def img2img(
     image = pipe(
         prompt,
         image=input_img,
+        strength=strength,
         guidance_scale=guidance, num_inference_steps=steps,
         generator=torch.Generator("cuda").manual_seed(seed)
     ).images[0]