add saving/recovering tmp user data for faster processing

app.py

@@ -75,11 +75,57 @@ def set_model_from_name(inp_model_name):
     inp_model = inpainting_models[inp_model_name]
 
 
+def save_user_session(hr_image, hr_mask, lr_results, prompt, session_id=None):
+    if session_id == '':
+        session_id = str(uuid.uuid4())
+    
+    tmp_dir = Path(TMP_DIR)
+    session_dir = tmp_dir / session_id
+    session_dir.mkdir(exist_ok=True, parents=True)
+    
+    hr_image.save(session_dir / 'hr_image.png')
+    hr_mask.save(session_dir / 'hr_mask.png')
+
+    lr_results_dir = session_dir / 'lr_results'
+    if lr_results_dir.exists():
+        shutil.rmtree(lr_results_dir)
+    lr_results_dir.mkdir(parents=True)
+    for i, lr_result in enumerate(lr_results):
+        lr_result.save(lr_results_dir / f'{i}.png')
+
+    with open(session_dir / 'prompt.txt', 'w') as f:
+        f.write(prompt)
+    
+    return session_id
+
+
+def recover_user_session(session_id):
+    if session_id == '':
+        return None, None, []
+    
+    tmp_dir = Path(TMP_DIR)
+    session_dir = tmp_dir / session_id
+    lr_results_dir = session_dir / 'lr_results'
+
+    hr_image = Image.open(session_dir / 'hr_image.png')
+    hr_mask = Image.open(session_dir / 'hr_mask.png')
+  
+    lr_result_paths = list(lr_results_dir.glob('*.png'))
+    gallery = []
+    for lr_result_path in sorted(lr_result_paths):
+        gallery.append(Image.open(lr_result_path))
+
+    with open(session_dir / 'prompt.txt', "r") as f:
+        prompt = f.read()
+
+    return hr_image, hr_mask, gallery, prompt
+
+
 def rasg_run(
     use_painta, prompt, input, seed, eta,
     negative_prompt, positive_prompt, ddim_steps,
     guidance_scale=7.5,
-    batch_size=1
+    batch_size=1, session_id=''
 ):
     torch.cuda.empty_cache()
 
@@ -119,15 +165,18 @@ def rasg_run(
             dilation=12
         )
         blended_images.append(blended_image)
-        inpainted_images.append(inpainted_image.numpy()[0])
+        inpainted_images.append(inpainted_image.pil())
+
+    session_id = save_user_session(
+        input['image'], input['mask'], inpainted_images, prompt, session_id=session_id)
 
-    return blended_images, inpainted_images
+    return blended_images, session_id
 
 
 def sd_run(use_painta, prompt, input, seed, eta,
     negative_prompt, positive_prompt, ddim_steps,
     guidance_scale=7.5,
-    batch_size=1
+    batch_size=1, session_id=''
 ):
     torch.cuda.empty_cache()
 
@@ -167,32 +216,37 @@ def sd_run(use_painta, prompt, input, seed, eta,
             dilation=12
         )
         blended_images.append(blended_image)
-        inpainted_images.append(inpainted_image.numpy()[0])
+        inpainted_images.append(inpainted_image.pil())
 
-    return blended_images, inpainted_images
+    session_id = save_user_session(
+        input['image'], input['mask'], inpainted_images, prompt, session_id=session_id)
+    
+    return blended_images, session_id
 
 
 def upscale_run(
-    prompt, input, ddim_steps, seed, use_sam_mask, gallery, img_index,
+    ddim_steps, seed, use_sam_mask, session_id, img_index,
     negative_prompt='',
     positive_prompt=', high resolution professional photo'
 ):
+    hr_image, hr_mask, gallery, prompt = recover_user_session(session_id)
+
+    if len(gallery) == 0:
+        return Image.open('./assets/sr_info.png')
+
     torch.cuda.empty_cache()
 
     seed = int(seed)
     img_index = int(img_index)
- 
+
     img_index = 0 if img_index < 0 else img_index
     img_index = len(gallery) - 1 if img_index >= len(gallery) else img_index
-    img_info = gallery[img_index if img_index >= 0 else 0]
-    inpainted_image = image_from_url_text(img_info)
-    lr_image = IImage(inpainted_image)
-    hr_image = IImage(input['image']).resize(2048)
-    hr_mask = IImage(input['mask']).resize(2048)
+    inpainted_image = gallery[img_index if img_index >= 0 else 0]
+
     output_image = sr.run(
         sr_model,
         sam_predictor,
-        lr_image,
+        inpainted_image,
         hr_image,
         hr_mask,
         prompt=prompt + positive_prompt,
@@ -203,8 +257,8 @@ def upscale_run(
         seed=seed,
         use_sam_mask=use_sam_mask
     )
-    output_image.info = input['image'].info  # save metadata
-    return output_image, output_image
+
+    return output_image
 
 
 def switch_run(use_rasg, model_name, *args):
@@ -316,8 +370,7 @@ with gr.Blocks(css='style.css') as demo:
             [input, prompt, example_container]
         )
 
-    mock_output_gallery = gr.Gallery([], columns = 4, visible=False)
-    mock_hires = gr.Image(label = "__MHRO__", visible = False)
+    session_id = gr.Textbox(value='', visible=False)
     html_info = gr.HTML(elem_id=f'html_info', elem_classes="infotext")
 
     inpaint_btn.click(
@@ -334,25 +387,24 @@ with gr.Blocks(css='style.css') as demo:
             positive_prompt,
             ddim_steps,
             guidance_scale,
-            batch_size
+            batch_size,
+            session_id
         ], 
-        outputs=[output_gallery, mock_output_gallery], 
+        outputs=[output_gallery, session_id], 
         api_name="inpaint"
     )
     upscale_btn.click(
         fn=upscale_run, 
         inputs=[
-            prompt, 
-            input,
             ddim_steps,
             seed,
             use_sam_mask,
-            mock_output_gallery,
+            session_id,
             html_info
         ],
-        outputs=[hires_image, mock_hires], 
+        outputs=[hires_image], 
         api_name="upscale",
-        _js="function(a, b, c, d, e, f, g){ return [a, b, c, d, e, f, selected_gallery_index()] }",
+        _js="function(a, b, c, d, e){ return [a, b, c, d, selected_gallery_index()] }",
     )
 
 demo.queue(max_size=20)

lib/methods/sr.py

@@ -73,6 +73,11 @@ def run(
     negative_prompt = '',
     use_sam_mask = False
 ):
+    hr_image_info = hr_image.info
+    lr_image = IImage(lr_image)
+    hr_image = IImage(hr_image).resize(2048)
+    hr_mask = IImage(hr_mask).resize(2048)
+
     torch.manual_seed(seed)
     dtype = ddim.vae.encoder.conv_in.weight.dtype
     device = ddim.vae.encoder.conv_in.weight.device
@@ -143,6 +148,6 @@ def run(
             fake_img=hr_result,
             mask=hr_mask_orig.alpha().data[0]
         )
-        return Image.fromarray(hr_result)
-    else:
-        return Image.fromarray(hr_result)
+    hr_result = Image.fromarray(hr_result)
+    hr_result.info = hr_image_info  # save metadata
+    return hr_result