initial commit

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+import constants
+import utils
+
+PREDICTOR = None
+
+
+def inference(image: np.ndarray, text: str, center_crop: bool):
+    num_steps = 10
+    if not text.lower().startswith("remove the"):
+        raise gr.Error("Instruction should start with 'Remove the' !")
+
+    image = Image.fromarray(image)
+    cropped_image, image = utils.preprocess_image(image, center_crop=center_crop)
+
+    utils.seed_everything()
+    prediction = PREDICTOR.predict(image, text, num_steps)
+
+    print("Num steps:", num_steps)
+
+    return cropped_image, prediction
+
+
+if __name__ == "__main__":
+    utils.setup_environment()
+
+    if not PREDICTOR:
+        PREDICTOR = utils.get_predictor()
+
+    sample_image, sample_instruction, sample_step = constants.EXAMPLES[3]
+
+    gr.Interface(
+        fn=inference,
+        inputs=[
+            gr.Image(type="numpy", value=sample_image, label="Source Image").style(
+                height=256
+            ),
+            gr.Textbox(
+                label="Instruction",
+                lines=1,
+                value=sample_instruction,
+            ),
+            gr.Checkbox(value=True, label="Center Crop", interactive=False),
+        ],
+        outputs=[
+            gr.Image(type="pil", label="Cropped Image").style(height=256),
+            gr.Image(type="pil", label="Output Image").style(height=256),
+        ],
+        allow_flagging="never",
+        examples=constants.EXAMPLES,
+        cache_examples=True,
+        title=constants.TITLE,
+        description=constants.DESCRIPTION,
+    ).launch()

constants.py

@@ -0,0 +1,25 @@
+TITLE = "Inst-Inpaint: Instructing to Remove Objects with Diffusion Models"
+
+DESCRIPTION = """
+<p style='text-align: center'>
+    <a href='http://instinpaint.abyildirim.com' target='_blank'>Project Page</a> | 
+    <a href='https://arxiv.org/abs/2304.03246' target='_blank'>Paper</a> | 
+    <a href='https://github.com/abyildirim/inst-inpaint' target='_blank'>GitHub Repo</a> | 
+</p>
+<p style='text-align: center'>
+    This demo demonstrates the Inst-Inpaint's abilities for instruction-based image inpainting.
+</p>
+"""
+
+EXAMPLES = [
+    ["examples/kite-boy.png", "Remove the colorful kite", True],
+    ["examples/cat-car.jpg", "Remove the car", True],
+    ["examples/bus-tree.jpg", "Remove the bus", True],
+    ["examples/cups.webp", "Remove the cup at the left", True],
+    ["examples/woman-fantasy.jpg", "Remove the woman", True],
+    ["examples/clock.png", "Remove the round clock at the center", True],
+    ["examples/woman.png", "Remove the woman at the left", True],
+    ["examples/men.png", "Remove the man at the right", True],
+    ["examples/tree.png", "Remove the tree", True],
+    ["examples/birds.png", "Remove the bird at the right of the bird", True]
+]

requirements.txt

@@ -0,0 +1,13 @@
+-f https://download.pytorch.org/whl/torch_stable.html
+git+https://github.com/openai/CLIP.git
+torch==1.13.1+cpu
+torchvision==0.14.1+cpu
+pytorch-lightning==1.6.5
+taming-transformers-rom1504==0.0.6
+einops==0.6.0
+kornia==0.6.11
+transformers==4.27.4
+dill==0.3.6
+gradio==3.24.1
+gdown==4.7.1
+torchmetrics==0.11.4

utils.py

@@ -0,0 +1,78 @@
+import logging
+import os
+import random
+import tarfile
+from typing import Tuple
+
+import dill
+import gdown
+import numpy as np
+import torch
+from PIL import Image
+from torchvision.transforms import ToTensor
+
+logger = logging.getLogger(__file__)
+
+to_tensor = ToTensor()
+
+
+def preprocess_image(
+    image: Image, resize_shape: Tuple[int, int] = (256, 256), center_crop=True
+):
+    processed_image = image
+
+    if center_crop:
+        width, height = image.size
+        crop_size = min(width, height)
+
+        left = (width - crop_size) // 2
+        top = (height - crop_size) // 2
+        right = (width + crop_size) // 2
+        bottom = (height + crop_size) // 2
+
+        processed_image = image.crop((left, top, right, bottom))
+
+    processed_image = processed_image.resize(resize_shape)
+
+    image = to_tensor(processed_image)
+    image = image.unsqueeze(0) * 2 - 1
+
+    return processed_image, image
+
+
+def download_artifacts(output_path: str):
+    logger.error("Downloading the model artifacts...")
+    if not os.path.exists(output_path):
+        gdown.download(id=os.environ["GDRIVE_ID"], output=output_path, quiet=True)
+
+
+def extract_artifacts(path: str):
+    logger.error("Extracting the model artifacts...")
+    if not os.path.exists("model.pkl"):
+        with tarfile.open(path) as tar:
+            tar.extractall()
+
+
+def setup_environment():
+    os.environ["PYTHONPATH"] = os.getcwd()
+
+    artifacts_path = "artifacts.tar.gz"
+
+    download_artifacts(output_path=artifacts_path)
+
+    extract_artifacts(path=artifacts_path)
+
+
+def get_predictor():
+    logger.error("Loading the predictor...")
+    with open("model.pkl", "rb") as fp:
+        return dill.load(fp)
+
+
+def seed_everything(seed: int = 0):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True