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

pull/4/head
Guillermo Rodriguez 2023-01-17 19:01:31 -03:00
parent 5a365180ff
commit 8427165a76
No known key found for this signature in database
GPG Key ID: EC3AB66D5D83B392
3 changed files with 12 additions and 1 deletions

View File

@ -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

View File

@ -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)

View File

@ -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]