diff --git a/modules/api/api.py b/modules/api/api.py index 97ec7514e..6de2ffb3b 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -562,7 +562,27 @@ class Api: img2imgreq.init_images = None img2imgreq.mask = None - return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js()) + from modules.s3 import s3_client + from modules.env_to_yaml import get_env_var + import uuid + + b64_string = b64images[0].decode("utf-8") + image_bytes = base64.b64decode(b64_string) + + uid_int = uuid.uuid4().int + + filename = str(uid_int)[0:18] + folder = "ddcn_results" + + s3_client.put_object( + Bucket=get_env_var(key="S3_BUCKET"), + Key=f"{folder}/{filename}.jpg", + Body=image_bytes, + ContentType="image/jpeg" + ) + s3_url = f"{get_env_var('S3_ENDPOINT')}/{get_env_var('S3_BUCKET')}/{folder}/{filename}.jpg" + + return models.ImageToImageResponse(parameters=vars(img2imgreq), info=processed.js(), s3_url=s3_url) def extras_single_image_api(self, req: models.ExtrasSingleImageRequest): reqDict = setUpscalers(req) @@ -924,5 +944,4 @@ class Api: def stop_webui(request): shared.state.server_command = "stop" - return Response("Stopping.") - + return Response("Stopping.") \ No newline at end of file diff --git a/modules/api/models.py b/modules/api/models.py index 79c4cbbbb..16497e94d 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -137,9 +137,9 @@ class TextToImageResponse(BaseModel): info: str class ImageToImageResponse(BaseModel): - images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.") parameters: dict info: str + s3_url: str class ExtrasBaseRequest(BaseModel): resize_mode: Literal[0, 1] = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.") diff --git a/modules/env_to_yaml.py b/modules/env_to_yaml.py new file mode 100644 index 000000000..a17d741e1 --- /dev/null +++ b/modules/env_to_yaml.py @@ -0,0 +1,8 @@ +from dotenv import load_dotenv +import os + +load_dotenv() + +def get_env_var(key: str) -> str: + value = os.getenv(key) + return value \ No newline at end of file diff --git a/modules/img2img.py b/modules/img2img.py index 24f869f5c..179cc6a99 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -17,12 +17,14 @@ from modules.ui import plaintext_to_html import modules.scripts -def process_batch(p, input, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0, use_png_info=False, png_info_props=None, png_info_dir=None): +def process_batch(p, input, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0, use_png_info=False, + png_info_props=None, png_info_dir=None): output_dir = output_dir.strip() processing.fix_seed(p) if isinstance(input, str): - batch_images = list(shared.walk_files(input, allowed_extensions=(".png", ".jpg", ".jpeg", ".webp", ".tif", ".tiff"))) + batch_images = list( + shared.walk_files(input, allowed_extensions=(".png", ".jpg", ".jpeg", ".webp", ".tif", ".tiff"))) else: batch_images = [os.path.abspath(x.name) for x in input] @@ -50,7 +52,7 @@ def process_batch(p, input, output_dir, inpaint_mask_dir, args, to_scale=False, batch_results = None discard_further_results = False for i, image in enumerate(batch_images): - state.job = f"{i+1} out of {len(batch_images)}" + state.job = f"{i + 1} out of {len(batch_images)}" if state.skipped: state.skipped = False @@ -105,7 +107,8 @@ def process_batch(p, input, output_dir, inpaint_mask_dir, args, to_scale=False, parsed_parameters = {} p.prompt = prompt + (" " + parsed_parameters["Prompt"] if "Prompt" in parsed_parameters else "") - p.negative_prompt = negative_prompt + (" " + parsed_parameters["Negative prompt"] if "Negative prompt" in parsed_parameters else "") + p.negative_prompt = negative_prompt + ( + " " + parsed_parameters["Negative prompt"] if "Negative prompt" in parsed_parameters else "") p.seed = int(parsed_parameters.get("Seed", seed)) p.cfg_scale = float(parsed_parameters.get("CFG scale", cfg_scale)) p.sampler_name = parsed_parameters.get("Sampler", sampler_name) @@ -149,7 +152,15 @@ def process_batch(p, input, output_dir, inpaint_mask_dir, args, to_scale=False, return batch_results -def img2img(id_task: str, request: gr.Request, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, mask_blur: int, mask_alpha: float, inpainting_fill: int, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, denoising_strength: float, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, img2img_batch_use_png_info: bool, img2img_batch_png_info_props: list, img2img_batch_png_info_dir: str, img2img_batch_source_type: str, img2img_batch_upload: list, *args): +def img2img(id_task: str, request: gr.Request, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, + sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, + init_mask_inpaint, mask_blur: int, mask_alpha: float, inpainting_fill: int, n_iter: int, batch_size: int, + cfg_scale: float, image_cfg_scale: float, denoising_strength: float, selected_scale_tab: int, height: int, + width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, + inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, + img2img_batch_inpaint_mask_dir: str, override_settings_texts, img2img_batch_use_png_info: bool, + img2img_batch_png_info_props: list, img2img_batch_png_info_dir: str, img2img_batch_source_type: str, + img2img_batch_upload: list, *args): override_settings = create_override_settings_dict(override_settings_texts) is_batch = mode == 5 @@ -229,10 +240,17 @@ def img2img(id_task: str, request: gr.Request, mode: int, prompt: str, negative_ output_dir = "" inpaint_mask_dir = "" png_info_dir = img2img_batch_png_info_dir if not shared.cmd_opts.hide_ui_dir_config else "" - processed = process_batch(p, img2img_batch_upload, output_dir, inpaint_mask_dir, args, to_scale=selected_scale_tab == 1, scale_by=scale_by, use_png_info=img2img_batch_use_png_info, png_info_props=img2img_batch_png_info_props, png_info_dir=png_info_dir) - else: # "from dir" + processed = process_batch(p, img2img_batch_upload, output_dir, inpaint_mask_dir, args, + to_scale=selected_scale_tab == 1, scale_by=scale_by, + use_png_info=img2img_batch_use_png_info, + png_info_props=img2img_batch_png_info_props, png_info_dir=png_info_dir) + else: # "from dir" assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled" - processed = process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args, to_scale=selected_scale_tab == 1, scale_by=scale_by, use_png_info=img2img_batch_use_png_info, png_info_props=img2img_batch_png_info_props, png_info_dir=img2img_batch_png_info_dir) + processed = process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, + img2img_batch_inpaint_mask_dir, args, to_scale=selected_scale_tab == 1, + scale_by=scale_by, use_png_info=img2img_batch_use_png_info, + png_info_props=img2img_batch_png_info_props, + png_info_dir=img2img_batch_png_info_dir) if processed is None: processed = Processed(p, [], p.seed, "") @@ -250,4 +268,5 @@ def img2img(id_task: str, request: gr.Request, mode: int, prompt: str, negative_ if opts.do_not_show_images: processed.images = [] - return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments") + return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html( + processed.comments, classname="comments") diff --git a/modules/s3.py b/modules/s3.py new file mode 100644 index 000000000..4ce3aeb0d --- /dev/null +++ b/modules/s3.py @@ -0,0 +1,15 @@ +import boto3 +from botocore.client import Config + +from modules.env_to_yaml import get_env_var + +s3_client = boto3.client( + 's3', + endpoint_url=get_env_var(key="S3_ENDPOINT"), + aws_access_key_id=get_env_var(key="S3_ACCESS_KEY"), + aws_secret_access_key=get_env_var(key="S3_SECRET_ACCESS_KEY"), + region_name='ru-1', + config=Config(s3={'addressing_style': 'path'}) +) + + diff --git a/requirements.txt b/requirements.txt index 0d6bac600..e048612a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,7 @@ clean-fid diskcache einops facexlib +boto3==1.34.148 fastapi>=0.90.1 gradio==3.41.2 inflection diff --git a/requirements_npu.txt b/requirements_npu.txt index 5e6a43646..fc601cb51 100644 --- a/requirements_npu.txt +++ b/requirements_npu.txt @@ -2,3 +2,4 @@ cloudpickle decorator synr==0.5.0 tornado +boto3==1.34.148 diff --git a/requirements_versions.txt b/requirements_versions.txt index 0306ce94f..1e84bb0e7 100644 --- a/requirements_versions.txt +++ b/requirements_versions.txt @@ -33,3 +33,5 @@ torchsde==0.2.6 transformers==4.30.2 httpx==0.24.1 pillow-avif-plugin==1.4.3 +boto3==1.34.148 + diff --git a/webui-user.bat b/webui-user.bat index e5a257bef..7c4c3768a 100644 --- a/webui-user.bat +++ b/webui-user.bat @@ -3,6 +3,6 @@ set PYTHON= set GIT= set VENV_DIR= -set COMMANDLINE_ARGS= +set COMMANDLINE_ARGS=--api --no-half call webui.bat