From ba3a32ae470a371fdc59341abddc5304f22841c2 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 8 Nov 2024 09:49:46 -0500 Subject: [PATCH] add api override field Signed-off-by: Vladimir Mandic --- cli/api-txt2img.js | 19 +--------- modules/api/control.py | 2 + modules/api/generate.py | 4 ++ modules/api/models.py | 2 + modules/images.py | 84 +++++++++++++++++++++-------------------- 5 files changed, 53 insertions(+), 58 deletions(-) diff --git a/cli/api-txt2img.js b/cli/api-txt2img.js index 46d09b3a2..8d0e9f5d1 100755 --- a/cli/api-txt2img.js +++ b/cli/api-txt2img.js @@ -20,23 +20,6 @@ const sd_options = { cfg_scale: 6, width: 512, height: 512, - /* - // enable second pass - enable_hr: true, - // second pass: upscale - hr_upscaler: 'SCUNet GAN', - hr_scale: 2.0, - // second pass: hires - hr_force: true, - hr_second_pass_steps: 20, - hr_sampler_name: 'UniPC', - denoising_strength: 0.5, - // second pass: refiner - refiner_steps: 5, - refiner_start: 0.8, - refiner_prompt: '', - refiner_negative: '', - */ // api return options save_images: false, send_images: true, @@ -55,7 +38,7 @@ async function main() { const json = await res.json(); console.log('result:', json.info); for (const i in json.images) { // eslint-disable-line guard-for-in - const f = `/tmp/test-{${i}.jpg`; + const f = `/tmp/test-${i}.jpg`; fs.writeFileSync(f, atob(json.images[i]), 'binary'); console.log('image saved:', f); } diff --git a/modules/api/control.py b/modules/api/control.py index cf8916095..ffb000053 100644 --- a/modules/api/control.py +++ b/modules/api/control.py @@ -31,6 +31,7 @@ ReqControl = models.create_model_from_signature( {"key": "ip_adapter", "type": Optional[List[models.ItemIPAdapter]], "default": None, "exclude": True}, {"key": "face", "type": Optional[models.ItemFace], "default": None, "exclude": True}, {"key": "control", "type": Optional[List[ItemControl]], "default": [], "exclude": True}, + {"key": "extra", "type": Optional[dict], "default": {}, "exclude": True}, ] ) @@ -159,6 +160,7 @@ class APIControl(): output_processed = [] output_info = '' run.control_set({ 'do_not_save_grid': not req.save_images, 'do_not_save_samples': not req.save_images, **self.prepare_ip_adapter(req) }) + run.control_set(getattr(req, "extra", {})) res = run.control_run(**args) for item in res: if len(item) > 0 and (isinstance(item[0], list) or item[0] is None): # output_images diff --git a/modules/api/generate.py b/modules/api/generate.py index aeafa05a0..e22102057 100644 --- a/modules/api/generate.py +++ b/modules/api/generate.py @@ -106,6 +106,8 @@ class APIGenerate(): p.scripts = script_runner p.outpath_grids = shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids p.outpath_samples = shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples + for key, value in getattr(txt2imgreq, "extra", {}).items(): + setattr(p, key, value) shared.state.begin('API TXT', api=True) script_args = script.init_script_args(p, txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner) if selectable_scripts is not None: @@ -150,6 +152,8 @@ class APIGenerate(): p.scripts = script_runner p.outpath_grids = shared.opts.outdir_img2img_grids p.outpath_samples = shared.opts.outdir_img2img_samples + for key, value in getattr(img2imgreq, "extra", {}).items(): + setattr(p, key, value) shared.state.begin('API-IMG', api=True) script_args = script.init_script_args(p, img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner) if selectable_scripts is not None: diff --git a/modules/api/models.py b/modules/api/models.py index 3cf3aade9..740f3c555 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -210,6 +210,7 @@ ReqTxt2Img = PydanticModelGenerator( {"key": "alwayson_scripts", "type": dict, "default": {}}, {"key": "ip_adapter", "type": Optional[List[ItemIPAdapter]], "default": None, "exclude": True}, {"key": "face", "type": Optional[ItemFace], "default": None, "exclude": True}, + {"key": "extra", "type": Optional[dict], "default": {}, "exclude": True}, ] ).generate_model() StableDiffusionTxt2ImgProcessingAPI = ReqTxt2Img @@ -235,6 +236,7 @@ ReqImg2Img = PydanticModelGenerator( {"key": "alwayson_scripts", "type": dict, "default": {}}, {"key": "ip_adapter", "type": Optional[List[ItemIPAdapter]], "default": None, "exclude": True}, {"key": "face_id", "type": Optional[ItemFace], "default": None, "exclude": True}, + {"key": "extra", "type": Optional[dict], "default": {}, "exclude": True}, ] ).generate_model() StableDiffusionImg2ImgProcessingAPI = ReqImg2Img diff --git a/modules/images.py b/modules/images.py index fb9cc9652..910349bef 100644 --- a/modules/images.py +++ b/modules/images.py @@ -40,8 +40,6 @@ def atomically_save_image(): except Exception: shared.log.warning(f'Save: unknown image format: {extension}') image_format = 'JPEG' - if shared.opts.image_watermark_enabled or (shared.opts.image_watermark_position != 'none' and shared.opts.image_watermark_image != ''): - image = set_watermark(image, shared.opts.image_watermark) exifinfo = (exifinfo or "") if shared.opts.image_metadata else "" # additional metadata saved in files if shared.opts.save_txt and len(exifinfo) > 0: @@ -153,6 +151,11 @@ def save_image(image, info = image.info.get(pnginfo_section_name, '') if info is not None: pnginfo[pnginfo_section_name] = info + + wm_text = getattr(p, 'watermark_text', shared.opts.image_watermark) + wm_image = getattr(p, 'watermark_image', shared.opts.image_watermark_image) + image = set_watermark(image, wm_text, wm_image) + params = script_callbacks.ImageSaveParams(image, p, filename, pnginfo) params.filename = namegen.sanitize(filename) dirname = os.path.dirname(params.filename) @@ -369,45 +372,46 @@ def draw_overlay(im, text: str = '', y_offset: int = 0): return im -def set_watermark(image, watermark): - if shared.opts.image_watermark_position != 'none': # visible watermark - wm_image = None - try: - wm_image = Image.open(shared.opts.image_watermark_image) - if wm_image.mode != 'RGBA': - wm_image = wm_image.convert('RGBA') - except Exception as e: - shared.log.warning(f'Set image watermark: fn="{shared.opts.image_watermark_image}" {e}') - if wm_image is not None: - if shared.opts.image_watermark_position == 'top/left': - position = (0, 0) - elif shared.opts.image_watermark_position == 'top/right': - position = (image.width - wm_image.width, 0) - elif shared.opts.image_watermark_position == 'bottom/left': - position = (0, image.height - wm_image.height) - elif shared.opts.image_watermark_position == 'bottom/right': - position = (image.width - wm_image.width, image.height - wm_image.height) - elif shared.opts.image_watermark_position == 'center': - position = ((image.width - wm_image.width) // 2, (image.height - wm_image.height) // 2) - else: - position = (random.randint(0, image.width - wm_image.width), random.randint(0, image.height - wm_image.height)) +def set_watermark(image, wm_text: str = None, wm_image: Image.Image = None): + if shared.opts.image_watermark_position != 'none' and wm_image is not None: # visible watermark + if isinstance(wm_image, str): try: - for x in range(wm_image.width): - for y in range(wm_image.height): - rgba = wm_image.getpixel((x, y)) - orig = image.getpixel((x+position[0], y+position[1])) - # alpha blend - a = rgba[3] / 255 - r = int(rgba[0] * a + orig[0] * (1 - a)) - g = int(rgba[1] * a + orig[1] * (1 - a)) - b = int(rgba[2] * a + orig[2] * (1 - a)) - if not a == 0: - image.putpixel((x+position[0], y+position[1]), (r, g, b)) - shared.log.debug(f'Set image watermark: fn="{shared.opts.image_watermark_image}" image={wm_image} position={position}') + wm_image = Image.open(wm_image) except Exception as e: shared.log.warning(f'Set image watermark: image={wm_image} {e}') + return image + if isinstance(wm_image, Image.Image): + if wm_image.mode != 'RGBA': + wm_image = wm_image.convert('RGBA') + if shared.opts.image_watermark_position == 'top/left': + position = (0, 0) + elif shared.opts.image_watermark_position == 'top/right': + position = (image.width - wm_image.width, 0) + elif shared.opts.image_watermark_position == 'bottom/left': + position = (0, image.height - wm_image.height) + elif shared.opts.image_watermark_position == 'bottom/right': + position = (image.width - wm_image.width, image.height - wm_image.height) + elif shared.opts.image_watermark_position == 'center': + position = ((image.width - wm_image.width) // 2, (image.height - wm_image.height) // 2) + else: + position = (random.randint(0, image.width - wm_image.width), random.randint(0, image.height - wm_image.height)) + try: + for x in range(wm_image.width): + for y in range(wm_image.height): + rgba = wm_image.getpixel((x, y)) + orig = image.getpixel((x+position[0], y+position[1])) + # alpha blend + a = rgba[3] / 255 + r = int(rgba[0] * a + orig[0] * (1 - a)) + g = int(rgba[1] * a + orig[1] * (1 - a)) + b = int(rgba[2] * a + orig[2] * (1 - a)) + if not a == 0: + image.putpixel((x+position[0], y+position[1]), (r, g, b)) + shared.log.debug(f'Set image watermark: image={wm_image} position={position}') + except Exception as e: + shared.log.warning(f'Set image watermark: image={wm_image} {e}') - if shared.opts.image_watermark_enabled: # invisible watermark + if shared.opts.image_watermark_enabled and wm_text is not None: # invisible watermark from imwatermark import WatermarkEncoder wm_type = 'bytes' wm_method = 'dwtDctSvd' @@ -416,16 +420,16 @@ def set_watermark(image, watermark): info = image.info data = np.asarray(image) encoder = WatermarkEncoder() - text = f"{watermark:<{length}}"[:length] + text = f"{wm_text:<{length}}"[:length] bytearr = text.encode(encoding='ascii', errors='ignore') try: encoder.set_watermark(wm_type, bytearr) encoded = encoder.encode(data, wm_method) image = Image.fromarray(encoded) image.info = info - shared.log.debug(f'Set invisible watermark: {watermark} method={wm_method} bits={wm_length}') + shared.log.debug(f'Set invisible watermark: {wm_text} method={wm_method} bits={wm_length}') except Exception as e: - shared.log.warning(f'Set invisible watermark error: {watermark} method={wm_method} bits={wm_length} {e}') + shared.log.warning(f'Set invisible watermark error: {wm_text} method={wm_method} bits={wm_length} {e}') return image