img rename save test

app.py

@@ -3,7 +3,7 @@ import datetime
 
 import toml
 import gradio as gr
-from PIL import Image, PngImagePlugin
+import shutil
 from huggingface_hub import HfApi, snapshot_download
 
 from pnginfo import read_info_from_image, send_paras
@@ -97,7 +97,7 @@ def control_ui():
     inp.select(lambda: 'inp', inputs=None, outputs=selection)
     return gen_btn, [prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_image, info_extract, ref_str, i2i_image, i2i_str, i2i_noise, overlay, inp_img, selection], [save, rand_seed, reuse_seed, reuse_img_vibe, reuse_img_i2i, reuse_img_inp]
 
-def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_image, info_extract, ref_str, i2i_image, i2i_str, i2i_noise, overlay, inp_img, selection, save):
+def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_image, info_extract, ref_str, i2i_image, i2i_str, i2i_noise, overlay, inp_img, selection):
     global today_count
     set_token(os.environ.get('token'))
     img_data, payload = generate_novelai_image(
@@ -110,17 +110,6 @@ def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps
         return gr.Image(value=None), payload
     today_count = get_count() + 1
     img = image_from_bytes(img_data)
-    if save:
-        save_path = client_config['save_path']
-        today = datetime.date.today().strftime('%Y-%m-%d')
-        today_path = os.path.join(save_path, today)
-        os.makedirs(today_path, exist_ok=True)
-        filename = str(today_count).rjust(5, '0') + '-' + str(payload['parameters']['seed']) + '.png'
-        pnginfo_data = PngImagePlugin.PngInfo()
-        for k, v in img.info.items():
-            pnginfo_data.add_text(k, str(v))
-        img.save(os.path.join(today_path, filename), pnginfo=pnginfo_data)
-        api.upload_folder(folder_path=today_path, path_in_repo=today_path, repo_id="P01yH3dr0n/naimages", repo_type="dataset", token=os.environ.get("hf_token"))
     return img, payload
 
 def preview_ui():
@@ -129,6 +118,22 @@ def preview_ui():
         info = gr.JSON(value={}, label="生成信息")
     return image, info
 
+def rename_save_img(path, payload, save):
+    if path is None:
+        return None
+    default = os.path.basename(path)
+    filename = str(today_count).rjust(5, '0') + '-' + str(payload['parameters']['seed']) + '.png'
+    renamed_path = path.replace(default, filename)
+    os.replace(path, renamed_path)
+    if save:
+        save_path = client_config['save_path']
+        today = datetime.date.today().strftime('%Y-%m-%d')
+        today_path = os.path.join(save_path, today)
+        os.makedirs(today_path, exist_ok=True)
+        shutil.copy(renamed_path, os.path.join(today_path, filename))
+        api.upload_folder(folder_path=today_path, path_in_repo=today_path, repo_id="P01yH3dr0n/naimages", repo_type="dataset", token=os.environ.get("hf_token"))
+    return renamed_path
+
 def main_ui():
     with gr.Blocks() as page:
         with gr.Row(variant="panel"):
@@ -136,11 +141,12 @@ def main_ui():
                 gen_btn, paras, others = control_ui()
             with gr.Column():
                 image, info = preview_ui()
-    gen_btn.click(generate, paras + [others[0]], [image, info], concurrency_limit=1)
+    gen_btn.click(generate, paras, [image, info], concurrency_limit=1)
     others[2].click(lambda o, s: o if len(s) == 0 else s['parameters']['seed'], inputs=[paras[3], info], outputs=paras[3])
     others[3].click(lambda i: i, inputs=image, outputs=paras[14])
     others[4].click(lambda i: i, inputs=image, outputs=paras[17])
-    others[5].click(lambda i: gr.ImageEditor(value={'background': i, 'layers': [Image.new('RGBA', i.size)], 'composite': i}), inputs=image, outputs=paras[21])
+    others[5].click(lambda i: gr.ImageEditor(value=i), inputs=image, outputs=paras[21])
+    image.change(rename_save_img, inputs=[image, info, others[0]], outputs=image)
     return page, paras[:14]
 
 def util_ui():