Update app.py

app.py

@@ -259,7 +259,7 @@ def inference(input_prompt, input_category):
     seed_everything(opt.seed)
     
     tic = time.time()
-    config = OmegaConf.load(f"{opt.config}")
+    config = x.load(f"{opt.config}")
     model = load_model_from_config(config, f"{opt.sd_ckpt}")
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     model = model.to(device)
@@ -359,7 +359,9 @@ def inference(input_prompt, input_category):
                         # cv2.imwrite(os.path.join("demo/demo_mask.png"), done_image_mask)
                         
                         # torchvision.utils.save_image(annotation_pred, os.path.join("demo/demo_segresult.png"), normalize=True, scale_each=True)
-                        return x_sample, done_image_mask
+                        generated_image = x_sample
+                        generated_mask = done_image_mask
+                        return [generated_image, generated_mask]
 
 
 # def make_transparent_foreground(pic, mask):
@@ -419,14 +421,154 @@ def inference(input_prompt, input_category):
 # model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True)
 # model.eval()
 
-gr.Interface(
-    inference,
-    gr.inputs.Textbox(label='Prompt', default='a photo of a lion on a mountain top at sunset'),
-    gr.inputs.Textbox(label='category', default='lion'),
-    gr.outputs.Image(type="pil", label="Output"),
-    # title=title,
-    # description=description,
-    # article=article,
-    # examples=[['demis.jpg'], ['lifeifei.png']],
-    # enable_queue=True
-).launch(debug=False)
+# gr.Interface(
+#     inference,
+#     gr.inputs.Textbox(label='Prompt', default='a photo of a lion on a mountain top at sunset'),
+#     gr.inputs.Textbox(label='category', default='lion'),
+#     gr.outputs.Image(type="pil", label="Output"),
+#     # title=title,
+#     # description=description,
+#     # article=article,
+#     # examples=[['demis.jpg'], ['lifeifei.png']],
+#     # enable_queue=True
+# ).launch(debug=False)
+
+def main():
+
+    # def load_example(
+    #     steps: int,
+    #     randomize_seed: bool,
+    #     seed: int,
+    #     randomize_cfg: bool,
+    #     text_cfg_scale: float,
+    #     image_cfg_scale: float,
+    # ):
+    #     example_instruction = random.choice(example_instructions)
+    #     return [example_image, example_instruction] + generate(
+    #         example_image,
+    #         example_instruction,
+    #         steps,
+    #         randomize_seed,
+    #         seed,
+    #         randomize_cfg,
+    #         text_cfg_scale,
+    #         image_cfg_scale,
+    #     )
+
+    # def generate(
+    #     input_image: Image.Image,
+    #     instruction: str,
+    #     steps: int,
+    #     randomize_seed: bool,
+    #     seed: int,
+    #     randomize_cfg: bool,
+    #     text_cfg_scale: float,
+    #     image_cfg_scale: float,
+    # ):
+    #     seed = random.randint(0, 100000) if randomize_seed else seed
+    #     text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
+    #     image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
+
+    #     width, height = input_image.size
+    #     factor = 512 / max(width, height)
+    #     factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
+    #     width = int((width * factor) // 64) * 64
+    #     height = int((height * factor) // 64) * 64
+    #     input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+    #     if instruction == "":
+    #         return [input_image, seed]
+
+    #     generator = torch.manual_seed(seed)
+    #     edited_image = pipe(
+    #         instruction, image=input_image,
+    #         guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
+    #         num_inference_steps=steps, generator=generator,
+    #     ).images[0]
+    #     return [seed, text_cfg_scale, image_cfg_scale, edited_image]
+
+    # def reset():
+    #     return [0, "Randomize Seed", 1371, "Fix CFG", 7.5, 1.5, None]
+
+    with gr.Blocks() as demo:
+        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">
+   InstructPix2Pix: Learning to Follow Image Editing Instructions
+</h1>
+<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
+<br/>
+<a href="https://huggingface.co/spaces/timbrooks/instruct-pix2pix?duplicate=true">
+<img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+<p/>""")
+        with gr.Row():
+            # with gr.Column(scale=1, min_width=100):
+            #     load_button = gr.Button("Load Example")
+            # with gr.Column(scale=1, min_width=100):
+            #     reset_button = gr.Button("Reset")
+            with gr.Column(scale=3):
+                Prompt = gr.Textbox(lines=1, label="Prompt", interactive=True)
+            with gr.Column(scale=2):
+                Category = gr.Textbox(lines=1, label="Category", interactive=True)
+            with gr.Column(scale=1, min_width=100):
+                generate_button = gr.Button("Generate")
+
+        with gr.Row():
+            generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+            generated_mask = gr.Image(label=f"Generated Mask", type="pil", interactive=False)
+            generated_image.style(height=512, width=512)
+            generated_mask.style(height=512, width=512)
+
+        # with gr.Row():
+        #     steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)
+        #     randomize_seed = gr.Radio(
+        #         ["Fix Seed", "Randomize Seed"],
+        #         value="Randomize Seed",
+        #         type="index",
+        #         show_label=False,
+        #         interactive=True,
+        #     )
+        #     seed = gr.Number(value=1371, precision=0, label="Seed", interactive=True)
+        #     randomize_cfg = gr.Radio(
+        #         ["Fix CFG", "Randomize CFG"],
+        #         value="Fix CFG",
+        #         type="index",
+        #         show_label=False,
+        #         interactive=True,
+        #     )
+        #     text_cfg_scale = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
+        #     image_cfg_scale = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
+
+        gr.Markdown(help_text)
+
+        # load_button.click(
+        #     fn=load_example,
+        #     inputs=[
+        #         steps,
+        #         randomize_seed,
+        #         seed,
+        #         randomize_cfg,
+        #         text_cfg_scale,
+        #         image_cfg_scale,
+        #     ],
+        #     outputs=[input_image, instruction, seed, text_cfg_scale, image_cfg_scale, edited_image],
+        # )
+        generate_button.click(
+            fn=inference,
+            inputs=[
+                Prompt,
+                Category,
+            ],
+            outputs=[generated_image, generated_mask],
+        )
+        # reset_button.click(
+        #     fn=reset,
+        #     inputs=[],
+        #     outputs=[steps, randomize_seed, seed, randomize_cfg, text_cfg_scale, image_cfg_scale, edited_image],
+        # )
+
+    demo.queue(concurrency_count=1)
+    demo.launch(share=False)
+
+
+
+if __name__ == "__main__":
+    main()