From 5301172159d3c51d18f9fa827a1024ecea96de28 Mon Sep 17 00:00:00 2001
From: xie river <chuanxie@amazon.com>
Date: Thu, 30 Mar 2023 16:04:04 +0000
Subject: [PATCH] can upload train images to s3 via webui

---
 localizations/zh_CN.json |  5 +++++
 modules/shared.py        |  1 +
 modules/ui.py            | 24 +++++++++++++++++++++++-
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/localizations/zh_CN.json b/localizations/zh_CN.json
index 2f44c978aea..09bb57b644d 100644
--- a/localizations/zh_CN.json
+++ b/localizations/zh_CN.json
@@ -840,5 +840,10 @@
     "Save Preview(s) Frequency (Epochs)": "保存预览频率 (Epochs)",
     "A generic prompt used to generate a sample image to verify model fidelity.": "用于生成样本图像以验证模型保真度的通用提示。",
     "Job detail":"训练任务详情",
+    "S3 bucket name for uploading train images":"上传训练图片集的S3桶名",
+    "Output S3 folder":"S3文件夹目录",
+    "Upload Train Images to S3":"上传训练图片到S3",
+    "Error, please configure a S3 bucket at settings page first":"失败，请先到设置页面配置S3桶名",
+    "Upload":"上传",
     "--------": "--------"
 }
diff --git a/modules/shared.py b/modules/shared.py
index a7f056db353..a50a276430d 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -423,6 +423,7 @@ def refresh_sagemaker_endpoints(username):
 }))
 
 options_templates.update(options_section(('saving-paths', "Paths for saving"), {
+    "train_files_s3bucket":OptionInfo("","S3 bucket name for uploading train images",component_args=hide_dirs),
     "outdir_samples": OptionInfo("", "Output directory for images; if empty, defaults to three directories below", component_args=hide_dirs),
     "outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output directory for txt2img images', component_args=hide_dirs),
     "outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output directory for img2img images', component_args=hide_dirs),
diff --git a/modules/ui.py b/modules/ui.py
index 9258d556279..8dc7f6b59d8 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -11,7 +11,8 @@
 import time
 import traceback
 from functools import partial, reduce
-
+import boto3
+import datetime
 import gradio as gr
 import gradio.routes
 import gradio.utils
@@ -1463,6 +1464,27 @@ def update_orig(image, state):
 
         with gr.Row().style(equal_height=False):
             with gr.Tabs(elem_id="train_tabs"):
+                ## Begin add s3 images upload interface by River
+                s3 = boto3.client('s3')
+                def upload_to_s3(imgs):
+                    username = shared.username 
+                    timestamp = datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S') 
+                    bucket_name = opts.train_files_s3bucket
+                    if bucket_name == '':
+                        return 'Error, please configure a S3 bucket at settings page first'
+                    folder_name = f"train-images/{username}/{timestamp}"
+                    for i, img in enumerate(imgs):
+                        filename = img.name.split('/')[-1]
+                        object_name = f"{folder_name}/{filename}"
+                        s3.upload_file(img.name, bucket_name, object_name)
+                    return f"{len(imgs)} images uploaded to S3 folder: s3://{bucket_name}/{folder_name}"
+                
+                with gr.Tab(label="Upload Train Images to S3"):
+                    upload_files = gr.Files(label="Files")
+                    url_output = gr.Textbox(label="Output S3 folder")
+                    sub_btn = gr.Button("Upload")
+                    sub_btn.click(fn=upload_to_s3, inputs=upload_files, outputs=url_output)
+                ## End add s3 images upload interface by River
                 with gr.Tab(label="Train Embedding"):
                     gr.HTML(value="<p style='margin-bottom: 0.7em'>Train an embedding; you must specify a directory with a set of 1:1 ratio images <a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\" style=\"font-weight:bold;\">[wiki]</a></p>")