added more tools

README.md

@@ -9,4 +9,12 @@ app_file: app.py
 pinned: false
 ---
 
+```bash
+conda create -n editing-tools python=3.9 -y
+conda activate editing-tools
+conda install -c "nvidia/label/cuda-11.7.0" cuda-toolkit cuda
+pip install -r requirements.txt
+python app.py
+```
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,36 +1,10 @@
-import os
-import torch
 import gradio as gr
-from video_watermark_remover import *
-from video_converter import *
-from image_converter import *
-from image_watermark_remover import *
-from typing import List
-from pydantic import BaseModel
-from lama_cleaner.server import main
+from watermark_remover import convert_video_to_frames, remove_image_watermark, remove_video_watermark
+from video_converter import convert_video
+from image_converter import convert_image
+from image_editing import edit_image
+from image_inpainting import inpaint
 
-class FakeLamaArgs(BaseModel):
-    host: str = "0.0.0.0"
-    port: int = 5000
-    model: str = 'lama'
-    hf_access_token: str = ""
-    sd_disable_nsfw: bool = False
-    sd_cpu_textencoder: bool = True
-    sd_run_local: bool = False
-    sd_enable_xformers: bool = False
-    local_files_only: bool = False
-    cpu_offload: bool = False
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-    gui: bool = False
-    gui_size: List[int] = [1000, 1000]
-    input: str = ''
-    disable_model_switch: bool = True
-    debug: bool = False
-    no_half: bool = False
-    disable_nsfw: bool = False
-    enable_xformers: bool = True if torch.cuda.is_available() else False
-    model_dir: str = None
-    output_dir: str = None
 
 css = """
     #remove_btn {
@@ -51,18 +25,27 @@ css = """
     #convert_btn:hover {
         background: linear-gradient(#2bbbc3, #201d18);
     }
+    #button {
+        background: linear-gradient(#201d18, #2bbbc3);
+        font-weight: bold;
+        font-size: 18px;
+        color:white;
+    }
+    #button:hover {
+        background: linear-gradient(#2bbbc3, #201d18);
+    }
     footer {
         display: none !important;
     }
 """
 
-demo = gr.Blocks(css=css, title="Video Watermark Remover")
+demo = gr.Blocks(css=css, title="Editing Tools")
 with demo:
     with gr.Tab("Image Converter"):
         gr.Markdown("""
         # <center>🖼️ Image Converter</center>
         """)
-        image_format = ['jpg', 'jpeg', 'png', 'bmp', 'tiff', 'gif', 'webp', 'ico', 'heic', 'heiv', 'heif']
+        image_format = ['jpg', 'jpeg', 'png', 'bmp', 'tiff', 'gif', 'webp', 'ico']
         with gr.Row():
             with gr.Column():
                 input_image = gr.File(label="Upload an Image")
@@ -87,6 +70,32 @@ with demo:
         output_image_clean = gr.Image(label="Output Image", interactive=True)
 
         image_remove_btn.click(remove_image_watermark, inputs=[input_image_watermark], outputs=[output_image_clean])
+    
+    with gr.Tab("Image Editing"):
+        gr.Markdown("""
+        # <center>🖼️ Image Editing</center>
+        """)
+        input_editing_image = gr.Image(label="Upload an Image", type="pil", interactive=True)
+        image_editing_options = gr.Radio(["High Res", "Colorize", "Greyscale", "Remove Background"], label="Select Editing Option", interactive=True, value="High Resolution")
+        image_editing_btn = gr.Button("Submit", interactive=True, elem_id="button")
+        with gr.Row():
+            image_editing_output = gr.Image(label="Output Preview", interactive=False)
+            image_editing_file = gr.File(label="Download File", interactive=False)
+
+        image_editing_btn.click(edit_image, inputs=[input_editing_image, image_editing_options], outputs=[image_editing_output, image_editing_file])
+
+    with gr.Tab("Image Inpainting"):
+        gr.Markdown("""
+        # <center>🖼️ Image Inpainting</center>
+        """)
+        input_inpainting_image = gr.Image(label="Upload an Image", type="pil", interactive=True, tool="sketch")
+        input_inpainting_prompt = gr.Textbox(label="Prompt", interactive=True)
+        input_inpainting_btn = gr.Button("Submit", interactive=True, elem_id="button")
+        with gr.Row():
+            input_inpainting_output = gr.Image(label="Image Preview", interactive=False)
+            input_inpainting_file = gr.File(label="Download File", interactive=False)
+
+        input_inpainting_btn.click(inpaint, inputs=[input_inpainting_image, input_inpainting_prompt], outputs=[input_inpainting_output, input_inpainting_file])
 
     with gr.Tab("Video Converter"):
         gr.Markdown("""
@@ -118,7 +127,7 @@ with demo:
         
     with gr.Tab("Video Watermark Remover"):
         gr.Markdown("""
-        # <center>🎥 Video Watermark Remover (Slow)</center>
+        # <center>🎥 Video Watermark Remover</center>
         """)
         with gr.Row():
             with gr.Column():
@@ -136,14 +145,6 @@ with demo:
         with gr.Row():
             output_video = gr.File(label="Output Video", interactive=False)
         input_video.change(convert_video_to_frames, inputs=[input_video], outputs=[mask, remove_btn])
-        remove_btn.click(remove_watermark, inputs=[mask], outputs=[output_video, remove_btn])
-    
-
-# Change the code according to the error
-import threading
-
-thread = threading.Thread(target=main, kwargs={'args': FakeLamaArgs()})
-thread.daemon = True
-thread.start()
+        remove_btn.click(remove_video_watermark, inputs=[mask], outputs=[output_video, remove_btn])
 
-demo.launch(show_api=False)
+demo.launch(show_api=False, share=True)

image_editing.py

@@ -0,0 +1,65 @@
+import os
+import sys
+from pathlib import Path
+import torch
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from gfpgan import GFPGANer
+from realesrgan import RealESRGANer
+import numpy as np
+import cv2
+from PIL import Image
+from rembg import remove
+
+# DeOldify
+os.system("hub install deoldify==1.2.0")
+import paddlehub as hub
+hub.server_check()
+colorize_model = hub.Module(name='deoldify')
+
+highres_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+bg_upsampler = RealESRGANer(
+    scale=4,
+    model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth',
+    model=highres_model,
+    tile=400,
+    tile_pad=10,
+    pre_pad=0,
+    half=True
+    )
+
+upsampler =  GFPGANer(
+    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
+    upscale=4,
+    arch='clean',
+    channel_multiplier=2,
+    bg_upsampler=bg_upsampler,
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    )
+
+
+os.makedirs("deoldify", exist_ok=True)
+os.makedirs("gfpganOutput", exist_ok=True)
+os.makedirs("greyscale", exist_ok=True)
+os.makedirs("rembg", exist_ok=True)
+
+def restore_image(image):
+     _, _, output = upsampler.enhance(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR), has_aligned=False, only_center_face=False, paste_back=True)
+     image = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+     return image
+
+def edit_image(image, option):
+    tools = ["High Res", "Colorize", "Greyscale", "Remove Background"]
+    if option == tools[0]:
+        restore_image(image).save("gfpganOutput/output.png")
+        return './gfpganOutput/output.png', './gfpganOutput/output.png'
+    elif option == tools[1]:
+        image.convert("L").save("deoldify/input.png", "PNG", quality=80, optimize=True)
+        colorize_model.predict("deoldify/input.png")
+        return './output/DeOldify/'+Path('deoldify/input.png').stem+".png", './output/DeOldify/'+Path('deoldify/input.png').stem+".png"
+    
+    elif option == tools[2]:
+        image.convert('L').save("greyscale/output.png")
+        return './greyscale/output.png', './greyscale/output.png'
+    elif option == tools[3]:
+        remove(image).save("rembg/output.png")
+        return './rembg/output.png', './rembg/output.png'

image_inpainting.py

@@ -0,0 +1,26 @@
+import os
+import torch
+from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers import AutoencoderKL
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", 
+        torch_dtype=torch.float16, 
+        revision="fp16",
+        vae=AutoencoderKL.from_pretrained(
+            "stabilityai/sd-vae-ft-mse",
+            torch_dtype=torch.float16
+            ).to(device)
+        ).to(device)
+pipe.enable_xformers_memory_efficient_attention()
+
+os.makedirs("inpainting_output", exist_ok=True)
+
+def inpaint(inputs, prompt):
+    image = inputs["image"].resize((image.size[0] - image.size[0] % 64, image.size[1] - image.size[1] % 64), Image.ANTIALIAS)
+    mask = inputs["mask"].resize((mask.size[0] - mask.size[0] % 64, mask.size[1] - mask.size[1] % 64), Image.ANTIALIAS)
+    output = pipe(prompt=prompt, image=image, mask_image=mask, guidance_scale=7.5, height=image.size[1], width=image.size[0])
+    output.images[0].save(f"inpainting_output/output.png")
+    return output.images[0], "inpainting_output/output.png"

image_variations.py

@@ -0,0 +1,62 @@
+import os
+import cv2
+from PIL import Image
+import numpy as np
+from diffusers import AutoencoderKL
+from diffusers import UniPCMultistepScheduler
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+# Blip for Image Captioning
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained(
+            "Salesforce/blip-image-captioning-base", 
+            torch_dtype=torch.float16).to(device)
+
+# ControlNet for Image Variation Generation based on Canny Edge Detection
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base", 
+    controlnet=ControlNetModel.from_pretrained(
+                "thibaud/controlnet-sd21-canny-diffusers", 
+                torch_dtype=torch.float16),
+    torch_dtype=torch.float16, 
+    revision="fp16",
+    vae=AutoencoderKL.from_pretrained(
+            "stabilityai/sd-vae-ft-mse",
+            torch_dtype=torch.float16
+            ).to(device)
+).to(device)
+
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.enable_xformers_memory_efficient_attention()
+
+def pre_process_image(image):
+  image = np.array(image)
+  low_threshold = 100
+  high_threshold = 200
+  image = cv2.Canny(image, low_threshold, high_threshold)
+  image = image[:, :, None]
+  image = np.concatenate([image, image, image], axis=2)
+  return Image.fromarray(image)
+
+def image_variations(image, input_prompt):
+    canny_image = pre_process_image(image)
+    if input_prompt:
+        prompt = input_prompt
+    else:
+        inputs = processor(image, return_tensors="pt").to(device, torch.float16)
+        out = model.generate(**inputs)
+        prompt = processor.decode(out[0], skip_special_tokens=True)
+        print(f"Blip Captioning: {prompt}")
+        
+    output_images = pipe(
+        [prompt]*4,
+        canny_image,
+        negative_prompt=["distorted, noisy, lowres, bad anatomy, worst quality, low quality, bad eyes, rough face, unclear face"] * 4,
+        num_inference_steps=25,
+    ).images
+
+    return output_images, canny_image

image_watermark_remover.py

@@ -1,61 +0,0 @@
-import io
-import requests
-from PIL import Image
-
-def remove_image_watermark(input):
-    image = input["image"].convert("RGB")
-    mask = input["mask"].convert("RGB")
-    image_data = io.BytesIO()
-    image.save(image_data, format="JPEG")
-    image_data = image_data.getvalue()
-
-    mask_data = io.BytesIO()
-    mask.save(mask_data, format="JPEG")
-    mask_data = mask_data.getvalue()
-
-    # Prepare form data
-    form_data = {
-            'ldmSteps': 25,
-            'ldmSampler': 'plms',
-            'zitsWireframe': True,
-            'hdStrategy': 'Original',
-            'hdStrategyCropMargin': 196,
-            'hdStrategyCropTrigerSize': 1280,
-            'hdStrategyResizeLimit': 2048,
-            'prompt': '',
-            'negativePrompt': '',
-            'croperX': -24,
-            'croperY': -23,
-            'croperHeight': 512,
-            'croperWidth': 512,
-            'useCroper': False,
-            'sdMaskBlur': 5,
-            'sdStrength': 0.75,
-            'sdSteps': 50,
-            'sdGuidanceScale': 7.5,
-            'sdSampler': 'pndm',
-            'sdSeed': 42,
-            'sdMatchHistograms': False,
-            'sdScale': 1,
-            'cv2Radius': 5,
-            'cv2Flag': 'INPAINT_NS',
-            'paintByExampleSteps': 50,
-            'paintByExampleGuidanceScale': 7.5,
-            'paintByExampleSeed': 42,
-            'paintByExampleMaskBlur': 5,
-            'paintByExampleMatchHistograms': False,
-            'sizeLimit': 1024,
-        }
-
-    files_data = {
-        'image': ('image.jpg', image_data),
-        'mask': ('mask.jpg', mask_data),
-    }
-
-    response = requests.post('http://localhost:5000/inpaint', data=form_data, files=files_data)
-
-    if response.headers['Content-Type'] == 'image/jpeg' or response.headers['Content-Type'] == 'image/png':
-        image = Image.open(io.BytesIO(response.content))
-        return image
-    else:
-        print(f"Error processing Image: {response.text}")

requirements.txt

@@ -1,8 +1,26 @@
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==1.13.1+cu117
+torchvision==0.14.1+cu117
+torchaudio==0.13.1
 gradio==3.22.1
 ffmpeg-python
 moviepy
 pydub
 opencv-python
 pyheif
-lama-cleaner==0.33.0
-xformers==0.0.16
+lama-cleaner==0.37.1
+xformers==0.0.16
+basicsr
+gfpgan
+realesrgan
+paddlepaddle-gpu
+imageio==2.4.1
+paddlehub
+triton==2.0.0
+git+https://github.com/muhammad-ahmed-ghani/rembg#egg=rembg[gpu]
+transformers
+ftfy
+scipy
+diffusers
+python-dotenv
+onnxruntime-gpu

video_converter.py

@@ -47,6 +47,14 @@ class VideoConverter:
         except Exception as e:
             raise Exception(f"Error converting audio: {e}")
 
+    def convert_to_gif(self, output_file):
+        try:
+            self.video.write_gif(output_file)
+            print(f"Video converted to GIF successfully!")
+            return output_file
+        except Exception as e:
+            raise Exception(f"Error converting video to GIF: {e}")
+
 def convert_video(input_file, format):
     try:
         converter = VideoConverter(input_file)
@@ -54,6 +62,8 @@ def convert_video(input_file, format):
             return converter.convert_video(f"output.{format}", format), "Converted video successfully!"
         elif format in ['mp3', 'wav', 'ogg', 'flac', 'aac']:
             return converter.convert_audio(f"output.{format}", format), "Converted audio successfully!"
+        elif format == "gif":
+            return converter.convert_to_gif(f"output.{format}"), "Converted to GIF successfully"
         else:
             return None, "Unsupported format!"
     except Exception as e:

video_watermark_remover.py → watermark_remover.py

@@ -1,94 +1,109 @@
 import glob
 import os
-import io
-import ffmpeg
-import requests
 from PIL import Image
 import shutil
 import concurrent.futures
 import gradio as gr
 import cv2
 import re
+import numpy as np
+import torch
+from lama_cleaner.helper import (
+    norm_img,
+    get_cache_path_by_url,
+    load_jit_model,
+)
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.schema import Config
+
+LAMA_MODEL_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+LAMA_MODEL_MD5 = os.environ.get("LAMA_MODEL_MD5", "e3aa4aaa15225a33ec84f9f4bc47e500")
+
+class LaMa(InpaintModel):
+    name = "lama"
+    pad_mod = 8
+
+    def init_model(self, device, **kwargs):
+        self.model = load_jit_model(LAMA_MODEL_URL, device, LAMA_MODEL_MD5).eval()
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        return os.path.exists(get_cache_path_by_url(LAMA_MODEL_URL))
+
+    def forward(self, image, mask, config: Config):
+        """Input image and output image have same size
+        image: [H, W, C] RGB
+        mask: [H, W]
+        return: BGR IMAGE
+        """
+        image = norm_img(image)
+        mask = norm_img(mask)
+
+        mask = (mask > 0) * 1
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+
+        inpainted_image = self.model(image, mask)
+
+        cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()
+        cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
+        cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
+        return cur_res
+
+lama_model = LaMa("cuda" if torch.cuda.is_available() else "cpu")
+config = Config(hd_strategy_crop_margin=196, ldm_steps=25, hd_strategy='Original', hd_strategy_crop_trigger_size=1280, hd_strategy_resize_limit=2048)
+
+def remove_image_watermark(inputs):
+    alpha_channel = None
+    image, mask = inputs["image"], inputs["mask"]
+    if image.mode == "RGBA":
+        image = np.array(image)
+        alpha_channel = image[:, :, -1]
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+    else:
+        image = np.array(image)
+    mask = cv2.threshold(np.array(mask.convert("L")), 127, 255, cv2.THRESH_BINARY)[1]
+    output = lama_model(image, mask, config)
+    output = cv2.cvtColor(output.astype(np.uint8), cv2.COLOR_BGR2RGB)
+    if alpha_channel is not None:
+        if alpha_channel.shape[:2] != output.shape[:2]:
+            alpha_channel = cv2.resize(
+                alpha_channel, dsize=(output.shape[1], output.shape[0])
+            )
+        output = np.concatenate(
+            (output, alpha_channel[:, :, np.newaxis]), axis=-1
+        )
+    return Image.fromarray(output)
 
 def process_image(mask_data, image_path):
-    image = Image.open(image_path)
-    image_data = io.BytesIO()
-    image.save(image_data, format=image.format)
-    image_data = image_data.getvalue()
-
-    # Prepare form data
-    form_data = {
-            'ldmSteps': 25,
-            'ldmSampler': 'plms',
-            'zitsWireframe': True,
-            'hdStrategy': 'Original',
-            'hdStrategyCropMargin': 196,
-            'hdStrategyCropTrigerSize': 1280,
-            'hdStrategyResizeLimit': 2048,
-            'prompt': '',
-            'negativePrompt': '',
-            'croperX': -24,
-            'croperY': -23,
-            'croperHeight': 512,
-            'croperWidth': 512,
-            'useCroper': False,
-            'sdMaskBlur': 5,
-            'sdStrength': 0.75,
-            'sdSteps': 50,
-            'sdGuidanceScale': 7.5,
-            'sdSampler': 'pndm',
-            'sdSeed': 42,
-            'sdMatchHistograms': False,
-            'sdScale': 1,
-            'cv2Radius': 5,
-            'cv2Flag': 'INPAINT_NS',
-            'paintByExampleSteps': 50,
-            'paintByExampleGuidanceScale': 7.5,
-            'paintByExampleSeed': 42,
-            'paintByExampleMaskBlur': 5,
-            'paintByExampleMatchHistograms': False,
-            'sizeLimit': 1024,
-        }
-
-    files_data = {
-        'image': (os.path.basename(image_path), image_data),
-        'mask': ('mask.png', mask_data)
-    }
-
-    response = requests.post('http://localhost:5000/inpaint', data=form_data, files=files_data)
-
-    if response.headers['Content-Type'] == 'image/jpeg' or response.headers['Content-Type'] == 'image/png':
-        output_image_path = os.path.join('output_images', os.path.splitext(os.path.basename(image_path))[0] + '_inpainted' + os.path.splitext(image_path)[1])
-        with open(output_image_path, 'wb') as output_image_file:
-            output_image_file.write(response.content)
-    else:
-        print(f"Error processing {image_path}: {response.text}")
+    output = remove_image_watermark({"image": Image.open(image_path), "mask": mask_data})
+    output_image_path = os.path.join('output_images', os.path.splitext(os.path.basename(image_path))[0] + '_inpainted' + os.path.splitext(image_path)[1])
+    output.save(output_image_path)
+    return output_image_path
 
-def remove_watermark(sketch, images_path='frames', output_path='output_images'):
+def remove_video_watermark(sketch, images_path='frames', output_path='output_images'):
     if os.path.exists('output_images'):
         shutil.rmtree('output_images')
     os.makedirs('output_images')
 
-    mask_data = io.BytesIO()
-    sketch["mask"].save(mask_data, format=sketch["mask"].format)
-    mask_data = mask_data.getvalue()
-
     image_paths = glob.glob(f'{images_path}/*.*')
 
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        executor.map(lambda image_path: process_image(mask_data, image_path), image_paths)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        executor.map(lambda image_path: process_image(sketch["mask"], image_path), image_paths)
 
     return gr.File.update(value=convert_frames_to_video('output_images'), visible=True), gr.Button.update(value='Done!')
 
 def convert_video_to_frames(video):
     if os.path.exists('input_video.mp4'):
         os.remove('input_video.mp4')
-    
+    # save the video to the current directory from temporary file
     with open(video, 'rb') as f:
         with open('input_video.mp4', 'wb') as f2:
             f2.write(f.read())
-            
-    #os.system(f"ffmpeg -i {video} input_video.mp4")
+    # os.system(f"ffmpeg -i {video} input_video.mp4")
     video_path = 'input_video.mp4'
 
     if os.path.exists('frames'):