22

app.py

@@ -89,8 +89,8 @@ def prepare(prompt, is_polish_prompt):
 def inference(
     prompt,
     negative_prompt,
-    input_image,
-    image_scale=1.0,
+    edit_dict,
+    # image_scale=1.0,
     control_mode='Canny',
     control_context_scale = 0.75,
     seed=42,
@@ -104,10 +104,11 @@ def inference(
 
     # process image
     print("DEBUG: process image")
-    if input_image is None:
-        print("Error: input_image is empty.")
+    if edit_dict is None:
+        print("Error: edit_dict is empty.")
         return None
     
+    print(edit_dict)
     # input_image, width, height = scale_image(input_image, image_scale)
     # control_mode='HED'
     processor_id = 'canny'
@@ -124,27 +125,29 @@ def inference(
     processor = Processor(processor_id)
 
     # Width must be divisible by 16
-    control_image, width, height = image_utils.rescale_image(input_image, image_scale, 16)
-    control_image = control_image.resize((1024, 1024))
+    
+    # control_image, width, height = image_utils.rescale_image(input_image, image_scale, 16)
+    # control_image = control_image.resize((1024, 1024))
+    width, height = edit_dict['background'].size
 
     print("DEBUG: control_image_torch")
     sample_size = [height, width]
-    control_image = processor(control_image, to_pil=True)
-    control_image = control_image.resize((width, height))
-    control_image_torch = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
+    # control_image = processor(control_image, to_pil=True)
+    # control_image = control_image.resize((width, height))
+    # control_image_torch = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
 
-    mask_image = None
-    # inpaint_image = None
-    
+    mask_image = edit_dict['composite']
     if mask_image is not None:
         mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
     else:
         mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
     
-    # if inpaint_image is not None:
-    #     inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
-    # else:
-    #     inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
+    inpaint_image = edit_dict['background']
+    if inpaint_image is not None:
+        inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
+    else:
+        inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
+    
 
     # generation
     if randomize_seed: seed = random.randint(0, MAX_SEED)
@@ -157,9 +160,9 @@ def inference(
         width=width,
         generator=generator,
         guidance_scale=guidance_scale,
-        image               = None,
+        image               = inpaint_image,
         mask_image          = mask_image,
-        control_image=control_image_torch,
+        # control_image=control_image_torch,
         num_inference_steps=num_inference_steps,
         control_context_scale=control_context_scale,
     ).images[0]
@@ -188,18 +191,12 @@ with gr.Blocks() as demo:
             gr.HTML(read_file("static/header.html"))
         with gr.Row():
             with gr.Column():
-                mask_image = gr.ImageEditor(
-                    height=290, 
-                    sources=['upload', 'clipboard'], 
-                    image_mode='RGB', 
-                    type="pil", label="Mask Image"
-                )
-
-                input_image = gr.Image(
+                edit_dict = gr.ImageEditor(
                     height=290, 
                     sources=['upload', 'clipboard'], 
                     image_mode='RGB', 
-                    type="pil", label="Upload"
+                    type="PIL",
+                    label="Mask Image"
                 )
                 
                 prompt = gr.Textbox(
@@ -240,23 +237,23 @@ with gr.Blocks() as demo:
                             step=0.01,
                             value=0.75,
                         )
-
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
+                    guidance_scale = gr.Slider(
                             label="Guidance scale",
                             minimum=0.0,
                             maximum=10.0,
                             step=0.1,
                             value=1.0,
                         )
+                    # with gr.Row():
+                        
 
-                        image_scale = gr.Slider(
-                            label="Image scale",
-                            minimum=0.5,
-                            maximum=2.0,
-                            step=0.1,
-                            value=1.0,
-                        )
+                        # image_scale = gr.Slider(
+                        #     label="Image scale",
+                        #     minimum=0.5,
+                        #     maximum=2.0,
+                        #     step=0.1,
+                        #     value=1.0,
+                        # )
 
                     seed = gr.Slider(
                         label="Seed",
@@ -274,10 +271,9 @@ with gr.Blocks() as demo:
                 with gr.Accordion("Preprocessor output", open=False):
                     control_image = gr.Image(label="Control image", show_label=False)
                     
-        gr.Examples(examples=examples, inputs=[input_image, prompt, control_mode])
+        gr.Examples(examples=examples, inputs=[edit_dict, prompt, control_mode])
         gr.Markdown(read_file("static/footer.md"))
 
-    mask_image.upload(fn=lambda x: x, inputs=[mask_image], outputs=[input_image])
     run_button.click(
         fn=prepare,
         inputs=[prompt, is_polish_prompt],
@@ -288,8 +284,8 @@ with gr.Blocks() as demo:
         inputs=[
             polished_prompt,
             negative_prompt,
-            input_image,
-            image_scale,
+            edit_dict,
+            # image_scale,
             control_mode,
             control_context_scale,
             seed,