Update app.py

app.py

@@ -51,13 +51,6 @@ def load_model_from_config(config, ckpt, verbose=False):
     pl_sd = torch.load(ckpt, map_location="cpu")
     sd = pl_sd["state_dict"]
     model = instantiate_from_config(config.model)
-    # m, u = model.load_state_dict(sd, strict=False)
-    # if len(m) > 0 and verbose:
-        # print("missing keys:")
-        # print(m)
-    # if len(u) > 0 and verbose:
-        # print("unexpected keys:")
-        # print(u)
     model.to(device)
     model.eval()
     return model
@@ -280,7 +273,6 @@ def inference(input_prompt, input_category):
                     data = [batch_size * [prompt]]
 
                 else:
-                    # print(f"reading prompts from {opt.from_file}")
                     with open(opt.from_file, "r") as f:
                         data = f.read().splitlines()
                         data = list(chunk(data, batch_size))
@@ -290,7 +282,6 @@ def inference(input_prompt, input_category):
 
                 start_code = None
                 if opt.fixed_code:
-                    # print('start_code')
                     start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
                 for n in trange(opt.n_iter, desc="Sampling"):
                     for prompts in tqdm(data, desc="data"):
@@ -320,7 +311,6 @@ def inference(input_prompt, input_category):
                         x_sample = torch.clamp((x_samples_ddim[0] + 1.0) / 2.0, min=0.0, max=1.0)
                         x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                         
-                        # Image.fromarray(x_sample.astype(np.uint8)).save("demo/demo.png")    
                         img = x_sample.astype(np.uint8)
                         
                         class_name = trainclass
@@ -352,154 +342,23 @@ def inference(input_prompt, input_category):
                         mask = annotation_pred.numpy()
                         mask = np.expand_dims(mask, 0)
                         done_image_mask = plot_mask(img, mask, alpha=0.9, indexlist=[0])
-                        # 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)
                         generated_image = x_sample.astype(np.uint8)
                         generated_mask = done_image_mask
                         return [generated_image, generated_mask]
 
 
-# def make_transparent_foreground(pic, mask):
-#     # split the image into channels
-#     b, g, r = cv2.split(np.array(pic).astype('uint8'))
-#     # add an alpha channel with and fill all with transparent pixels (max 255)
-#     a = np.ones(mask.shape, dtype='uint8') * 255
-#     # merge the alpha channel back
-#     alpha_im = cv2.merge([b, g, r, a], 4)
-#     # create a transparent background
-#     bg = np.zeros(alpha_im.shape)
-#     # setup the new mask
-#     new_mask = np.stack([mask, mask, mask, mask], axis=2)
-#     # copy only the foreground color pixels from the original image where mask is set
-#     foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8)
-
-#     return foreground
-
-
-# def remove_background(input_image):
-#     preprocess = transforms.Compose([
-#         transforms.ToTensor(),
-#         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-#     ])
-
-#     input_tensor = preprocess(input_image)
-#     input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model
-
-#     # move the input and model to GPU for speed if available
-#     if torch.cuda.is_available():
-#         input_batch = input_batch.to('cuda')
-#         model.to('cuda')
-
-#     with torch.no_grad():
-#         output = model(input_batch)['out'][0]
-#     output_predictions = output.argmax(0)
-
-#     # create a binary (black and white) mask of the profile foreground
-#     mask = output_predictions.byte().cpu().numpy()
-#     background = np.zeros(mask.shape)
-#     bin_mask = np.where(mask, 255, background).astype(np.uint8)
-
-#     foreground = make_transparent_foreground(input_image, bin_mask)
-
-#     return foreground, bin_mask
-
-
-# def inference(img):
-#     foreground, _ = remove_background(img)
-#     return foreground
-
-
-# torch.hub.download_url_to_file('https://pbs.twimg.com/profile_images/691700243809718272/z7XZUARB_400x400.jpg',
-#                                'demis.jpg')
-# torch.hub.download_url_to_file('https://hai.stanford.edu/sites/default/files/styles/person_medium/public/2020-03/hai_1512feifei.png?itok=INFuLABp',
-#                                'lifeifei.png')
-# 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)
-
 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
+   Guiding Text-to-Image Diffusion Model Towards Grounded Generation
 </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">
+<a href="https://huggingface.co/spaces/Purple11/Grounded-Diffusion?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):
@@ -513,40 +372,7 @@ def main():
             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=[
@@ -555,11 +381,6 @@ def main():
             ],
             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)