Add image resize before img2img calc, add strength parameter to local img2img tools

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

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)

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]