disable tensorflow library

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-import io
+import spaces
 from omegaconf import OmegaConf
 
 import subprocess
@@ -14,6 +14,7 @@ sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lama')
 
 from lama.saicinpainting.evaluation.refinement import refine_predict
 from lama.saicinpainting.training.trainers import load_checkpoint
+from lama.saicinpainting.evaluation.utils import move_to_device
 
 
 # Load the model
@@ -43,25 +44,20 @@ def get_inpaint_model():
     model.to(device)
     return model, predict_config
 
-def inpaint(input_dict):
+@spaces.GPU
+def inpaint(input_dict, refinement_enabled=False):
     """
     Performs image inpainting on the input image using the provided mask.
     Args: input_dict containing 'background' (image) and 'layers' (mask)
     Returns: Tuple of (output_image, input_mask)
     """
-    input_image = input_dict["background"].convert("RGB")
+    input_image = np.array(input_dict["background"].convert("RGB")).astype('float32') / 255
     input_mask = pil_to_binary_mask(input_dict['layers'][0])
 
-    # TODO: check if this is correct; (C,H,W) or (H,W,C)
-
-    # batch = dict(image=input_image, mask=input_mask[None, ...])
     np_input_image = np.transpose(np.array(input_image), (2, 0, 1))
     np_input_mask = np.array(input_mask)[None, :, :]  # Add channel dimension for grayscale images
     batch = dict(image=np_input_image, mask=np_input_mask)
 
-    print('lol', batch['image'].shape)
-    print('lol', batch['mask'].shape)
-
     inpaint_model, predict_config = get_inpaint_model()
     device = torch.device(predict_config.device)
 
@@ -69,8 +65,20 @@ def inpaint(input_dict):
     batch['image'] = torch.tensor(pad_img_to_modulo(batch['image'], predict_config.dataset.pad_out_to_modulo))[None].to(device)
     batch['mask'] = torch.tensor(pad_img_to_modulo(batch['mask'], predict_config.dataset.pad_out_to_modulo))[None].float().to(device)
 
-    cur_res = refine_predict(batch, inpaint_model, **predict_config.refiner)
-    cur_res = cur_res[0].permute(1,2,0).detach().cpu().numpy()
+
+    if refinement_enabled is True:
+        cur_res = refine_predict(batch, inpaint_model, **predict_config.refiner)
+        cur_res = cur_res[0].permute(1,2,0).detach().cpu().numpy()
+    else:
+        with torch.no_grad():
+            batch = move_to_device(batch, device)
+            batch['mask'] = (batch['mask'] > 0) * 1
+            batch = inpaint_model(batch)
+            cur_res = batch[predict_config.out_key][0].permute(1, 2, 0).detach().cpu().numpy()
+            unpad_to_size = batch.get('unpad_to_size', None)
+            if unpad_to_size is not None:
+                orig_height, orig_width = unpad_to_size
+                cur_res = cur_res[:orig_height, :orig_width]
 
     cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
     output_image = Image.fromarray(cur_res)
@@ -88,7 +96,7 @@ def pad_img_to_modulo(img, mod):
     out_width = ceil_modulo(width, mod)
     return np.pad(img, ((0, 0), (0, out_height - height), (0, out_width - width)), mode='symmetric')
 
-def pil_to_binary_mask(pil_image, threshold=0):
+def pil_to_binary_mask(pil_image, threshold=0, max_scale=1):
     """
     Converts a PIL image to a binary mask.
 
@@ -107,25 +115,35 @@ def pil_to_binary_mask(pil_image, threshold=0):
         for j in range(binary_mask.shape[1]):
             if binary_mask[i,j] == True :
                 mask[i,j] = 1
-    mask = (mask*255).astype(np.uint8)
+    mask = (mask*max_scale).astype(np.uint8)
     output_mask = Image.fromarray(mask)
     # Convert mask to grayscale
     return output_mask.convert("L")
 
+css = ".output-image, .input-image, .image-preview {height: 600px !important}"
+
 # Create Gradio interface
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     gr.Markdown("# Image Inpainting")
     gr.Markdown("Upload an image and draw a mask to remove unwanted objects.")
     
     with gr.Row():
-        input_image = gr.ImageEditor(type="pil", label='Input image & Mask', interactive=True, height="auto", width="auto")
-        output_image = gr.Image(type="pil", label="Output Image")
-    # with gr.Column():
-    #     masked_image = gr.Image(label="Masked image", type="pil")
+        input_image = gr.ImageEditor(type="pil", label='Input image & Mask', interactive=True, height="auto", width="auto", brush=gr.Brush(colors=['#f2e2cd'], default_size=25))
+        output_image = gr.Image(type="pil", label="Output Image", height="auto", width="auto")
+    
+    with gr.Row():
+        refine_checkbox = gr.Checkbox(label="Enable Refinement[SLOWER BUT BETTER]", value=False)
+        inpaint_button = gr.Button("Inpaint")
+
+    def inpaint_with_refinement(image, enable_refinement):
+        return inpaint(image, refinement_enabled=enable_refinement)
 
-    inpaint_button = gr.Button("Inpaint")
-    inpaint_button.click(fn=inpaint, inputs=[input_image], outputs=[output_image])
+    inpaint_button.click(
+        fn=inpaint_with_refinement,
+        inputs=[input_image, refine_checkbox],
+        outputs=[output_image]
+    )
 
 # Launch the interface
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -1,5 +1,5 @@
 gradio
-numpy
+numpy==1.26.4
 pillow
 pyyaml
 tqdm
@@ -7,7 +7,7 @@ easydict==1.9.0
 scikit-image
 scikit-learn
 opencv-python
-tensorflow
+# tensorflow
 joblib
 matplotlib
 pandas
@@ -21,3 +21,4 @@ packaging
 wldhx.yadisk-direct
 torch
 torchvision
+spaces

setup_local.sh

@@ -6,9 +6,19 @@ conda install pytorch torchvision -c pytorch -y
 pip install -r requirements.txt
 
 
-# Clone dependency repos
-git clone https://github.com/advimman/lama.git
+# Check if lama directory exists
+if [ ! -d "lama" ]; then
+    # Clone dependency repos
+    git clone https://github.com/advimman/lama.git
+else
+    echo "lama directory already exists. Skipping clone."
+fi
 
-# Download big-lama model
-curl -LJO https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip
-unzip big-lama.zip
+# Check if big-lama.zip or big-lama directory exists
+if [ ! -f "big-lama.zip" ] && [ ! -d "big-lama" ]; then
+    # Download big-lama model
+    curl -LJO https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip
+    unzip big-lama.zip
+else
+    echo "big-lama model already exists. Skipping download and extraction."
+fi