Initialize the repo

app.py

@@ -0,0 +1,178 @@
+import os
+import torch
+from diffusers import DDIMScheduler,DiffusionPipeline
+import torch.nn.functional as F
+import cv2
+from torchvision.utils import save_image
+from diffusers.utils import load_image
+from torchvision.transforms.functional import to_tensor, gaussian_blur
+from matplotlib import pyplot as plt
+import gradio as gr
+import spaces
+from gradio_imageslider import ImageSlider
+from torchvision.transforms.functional import to_pil_image, to_tensor
+from PIL import ImageFilter
+
+
+def preprocess_image(input_image, device):
+    image = to_tensor(input_image)
+    image = image.unsqueeze_(0).float() * 2 - 1 # [0,1] --> [-1,1]
+    if image.shape[1] != 3:
+        image = image.expand(-1, 3, -1, -1)
+    image = F.interpolate(image, (1024, 1024))
+    image = image.to(dtype).to(device)
+    
+    return image
+
+
+def preprocess_mask(input_mask, device):
+    mask = to_tensor(input_mask.convert('L'))
+    mask = mask.unsqueeze_(0).float()  # 0 or 1
+    mask = F.interpolate(mask, (1024, 1024))
+    mask = gaussian_blur(mask, kernel_size=(77, 77))
+    mask[mask < 0.1] = 0
+    mask[mask >= 0.1] = 1
+    mask = mask.to(dtype).to(device)
+    
+    return mask
+
+
+def make_redder(img, mask, increase_factor=0.4):
+    img_redder = img.clone()
+    mask_expanded = mask.expand_as(img)
+    img_redder[0][mask_expanded[0] == 1] = torch.clamp(img_redder[0][mask_expanded[0] == 1] + increase_factor, 0, 1)
+    
+    return img_redder
+
+
+# Model loading parameters
+is_cpu_offload_enabled = False
+is_attention_slicing_enabled = True
+
+# Load model
+dtype = torch.float16
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+
+model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+pipeline = DiffusionPipeline.from_pretrained(
+    model_path,
+    custom_pipeline="pipeline_stable_diffusion_xl_attentive_eraser.py",
+    scheduler=scheduler,
+    variant="fp16",
+    use_safetensors=True,
+    torch_dtype=dtype,
+).to(device)
+
+if is_attention_slicing_enabled:
+    pipeline.enable_attention_slicing()
+    
+if is_cpu_offload_enabled:
+    pipeline.enable_model_cpu_offload()
+
+
+@spaces.GPU
+def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, strength=0.8):
+    generator = torch.Generator('cuda').manual_seed(seed)
+    prompt = "" # Set prompt to null
+
+    source_image_pure = gradio_image["background"]
+    mask_image_pure = gradio_image["layers"][0]
+    source_image = preprocess_image(source_image_pure, device)
+    mask = preprocess_mask(mask_image_pure, device)
+    
+    START_STEP = 0 # AAS start step
+    END_STEP = int(strength * num_inference_steps) # AAS end step
+    LAYER = 34 # 0~23down,24~33mid,34~69up /AAS start layer 
+    END_LAYER = 70 # AAS end layer
+    ss_steps = 9 # similarity suppression steps
+    ss_scale = 0.3 # similarity suppression scale
+    
+    image = pipeline(
+        prompt=prompt,
+        image=source_image,
+        mask_image=mask,
+        height=1024,
+        width=1024,
+        AAS=True, # enable AAS
+        strength=strength, # inpainting strength
+        rm_guidance_scale=rm_guidance_scale, # removal guidance scale
+        ss_steps = ss_steps, # similarity suppression steps
+        ss_scale = ss_scale, # similarity suppression scale
+        AAS_start_step=START_STEP, # AAS start step
+        AAS_start_layer=LAYER, # AAS start layer 
+        AAS_end_layer=END_LAYER, # AAS end layer
+        num_inference_steps=num_inference_steps, # number of inference steps # AAS_end_step = int(strength*num_inference_steps)
+        generator=g,
+        guidance_scale=1,
+        output_type='pt'
+    ).images[0]
+    
+    img = (source_image * 0.5 + 0.5).squeeze(0)
+    mask_red = mask.squeeze(0)
+    img_redder = make_redder(img, mask_red)
+    
+    pil_mask = to_pil_image(mask.squeeze(0))
+    pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
+    mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
+    mask_f = 1-(1 - mask) * (1 - mask_blurred)
+    
+    image_1 = image.unsqueeze(0)
+    
+    return source_image, image_1
+    
+
+title = """<h1 align="center">Object Remove</h1>"""
+with gr.Blocks() as demo:
+    gr.HTML(title)
+    with gr.Row():
+        with gr.Column():
+            with gr.Accordion("Advanced Options", open=False):
+                guidance_scale = gr.Slider(
+                    minimum=1,
+                    maximum=20,
+                    value=9,
+                    step=0.1,
+                    label="Guidance Scale"
+                )
+                num_steps = gr.Slider(
+                    minimum=5,
+                    maximum=100,
+                    value=50,
+                    step=1,
+                    label="Steps"
+                )
+                seed = gr.Slider(
+                    minimum=42,
+                    maximum=100000000000,
+                    value=42,
+                    step=1,
+                    label="Seed"
+                )
+                strength = gr.Slider(
+                    minimum=0,
+                    maximum=1,
+                    value=0.8,
+                    step=0.1,
+                    label="Strength"
+                )
+            
+            input_image = gr.ImageMask(
+                type="pil", label="Input Image",crop_size=(1200,1200), layers=False
+            )
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    run_button = gr.Button("Generate")
+
+            result = ImageSlider(
+                interactive=False,
+                label="Generated Image",
+                type="pil"
+            )
+    
+    run_button.click(
+        fn=remove,
+        inputs=[input_image, guidance_scale, num_steps, seed, strength],
+        outputs=result,
+    )