no prompt button, api call to int

app.py

@@ -1,26 +1,45 @@
 import gradio as gr
-import torch
 import numpy as np
-import diffusers
+
 import os
-import spaces
-from PIL import Image
-hf_token = os.environ.get("HF_TOKEN")
-from diffusers import StableDiffusionXLInpaintPipeline, DDIMScheduler, UNet2DConditionModel
-from diffusers import (
-    AutoencoderKL,
-    LCMScheduler,
-)
-from pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline
-from controlnet import ControlNetModel, ControlNetConditioningEmbedding
-import torch
-import numpy as np
 from PIL import Image
 import requests
-import PIL
 from io import BytesIO
-from torchvision import transforms
+import io
+import base64
+
+hf_token = os.environ.get("HF_TOKEN")
+auth_headers = {"api_token": hf_token}
+
+def convert_mask_image_to_base64_string(mask_image):
+    buffer = io.BytesIO()
+    mask_image.save(buffer, format="PNG")  # You can choose the format (e.g., "JPEG", "PNG")
+    # Encode the buffer in base64
+    image_base64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    return f",{image_base64_string}" # for some reason the funciton which downloads image from base64 expects prefix of "," which is redundant in the url
+
+def download_image(url):
+    response = requests.get(url)
+    return Image.open(BytesIO(response.content)).convert("RGB")
 
+def eraser_api_call(image_base64_file, mask_base64_file, seed, mask_type, original_quality, guidance_scale):
+
+    # url = "http://engine.prod.bria-api.com/v1/eraser" # TODO: use this link!
+    url = "http://engine.int.bria-api.com/v1/eraser" # TODO: use this link!
+    
+    payload = {
+    "file": image_base64_file,
+    "mask_file": mask_base64_file,
+    "seed": seed,
+    "mask_type": mask_type,
+    "original_quality": original_quality,
+    "text_guidance_scale": guidance_scale
+    }
+    response = requests.post(url, json=payload, headers=auth_headers)
+    response = response.json()
+    res_image = download_image(response["result_url"])
+    
+    return res_image
 
 ratios_map =  {
     0.5:{"width":704,"height":1408},
@@ -43,13 +62,6 @@ ratios_map =  {
 }
 ratios = np.array(list(ratios_map.keys()))
 
-image_transforms = transforms.Compose(
-    [
-        transforms.ToTensor(),
-    ]
-)
-
-default_negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
 
 
 def get_masked_image(image, image_mask, width, height):
@@ -77,25 +89,6 @@ def get_size(init_image):
 
     return w,h
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load, init model
-controlnet = ControlNetModel().from_pretrained("briaai/DEV-ControlNetInpaintingFast", torch_dtype=torch.float16)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained("briaai/BRIA-2.3", controlnet=controlnet.to(dtype=torch.float16), torch_dtype=torch.float16, vae=vae) #force_zeros_for_empty_prompt=False, # vae=vae)
-
-pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights("briaai/BRIA-2.3-FAST-LORA")
-pipe.fuse_lora()
-
-pipe = pipe.to(device)
-# pipe.enable_xformers_memory_efficient_attention()
-
-# generator = torch.Generator(device='cuda').manual_seed(123456)
-
-vae = pipe.vae
-
-pipe.enable_model_cpu_offload() 
 
 def read_content(file_path: str) -> str:
     """read the content of target file
@@ -105,62 +98,19 @@ def read_content(file_path: str) -> str:
 
     return content
 
-@spaces.GPU(enable_queue=True)
-def predict(dict, prompt="", negative_prompt = default_negative_prompt, guidance_scale=1.2, steps=12, seed=123456):
-    if negative_prompt == "":
-        negative_prompt = None
+def predict(dict, guidance_scale=1.2, seed=123456):
 
     init_image = Image.fromarray(dict['background'][:, :, :3], 'RGB') #dict['background'].convert("RGB")#.resize((1024, 1024))
     mask = Image.fromarray(dict['layers'][0][:,:,3], 'L') #dict['layers'].convert("RGB")#.resize((1024, 1024))
-
-
-    width, height = get_size(init_image)
-
-    init_image = init_image.resize((width, height))
-    mask = mask.resize((width, height))
-    
-
-    masked_image, image_mask, masked_image_to_present = get_masked_image(init_image, mask, width, height)
-    masked_image_tensor = image_transforms(masked_image)
-    masked_image_tensor = (masked_image_tensor - 0.5) / 0.5
     
-    masked_image_tensor = masked_image_tensor.unsqueeze(0).to(device="cuda")
+    image_base64_file = convert_mask_image_to_base64_string(init_image)
+    mask_base64_file = convert_mask_image_to_base64_string(mask)
     
-    control_latents = vae.encode(  
-            masked_image_tensor[:, :3, :, :].to(vae.dtype)
-        ).latent_dist.sample()
+    mask_type = "brush"
+    original_quality = True
+    gen_img = eraser_api_call(image_base64_file, mask_base64_file, seed, mask_type, original_quality, guidance_scale)
     
-    control_latents = control_latents * vae.config.scaling_factor 
-    
-    image_mask = np.array(image_mask)[:,:]
-    mask_tensor = torch.tensor(image_mask, dtype=torch.float32)[None, ...]
-    # binarize the mask
-    mask_tensor = torch.where(mask_tensor > 128.0, 255.0, 0)       
-    
-    mask_tensor = mask_tensor / 255.0
-    
-    mask_tensor = mask_tensor.to(device="cuda")
-    mask_resized = torch.nn.functional.interpolate(mask_tensor[None, ...], size=(control_latents.shape[2], control_latents.shape[3]), mode='nearest')
-    # mask_resized = mask_resized.to(torch.float16)
-    masked_image = torch.cat([control_latents, mask_resized], dim=1)
-
-    generator = torch.Generator(device='cuda').manual_seed(int(seed))
-    
-    output = pipe(prompt = prompt,
-                  width=width,
-                  height=height,
-                  negative_prompt=negative_prompt,
-                  image = masked_image, # control image V
-                  init_image = init_image,
-                  mask_image = mask_tensor,
-                  guidance_scale = guidance_scale,
-                  num_inference_steps=int(steps),
-                  # strength=strength,
-                  generator=generator,
-                  controlnet_conditioning_sale=1.0)
-
-    torch.cuda.empty_cache
-    return output.images[0] #, gr.update(visible=True)
+    return gen_img
 
 
 css = '''
@@ -212,29 +162,22 @@ with image_blocks as demo:
           </p>
         ''')
     with gr.Row():
-                with gr.Column():
-                    image = gr.ImageEditor(sources=["upload"], layers=False, transforms=[], brush=gr.Brush(colors=["#000000"], color_mode="fixed")) #gr.Image(sources=['upload'], tool='sketch', elem_id="image_upload", type="pil", label="Upload", height=400)
-                    with gr.Row(elem_id="prompt-container", equal_height=True):
-                        with gr.Row():
-                            prompt = gr.Textbox(placeholder="Your prompt (what you want in place of what is erased)", show_label=False, elem_id="prompt")
-                            btn = gr.Button("Inpaint!", elem_id="run_button")
-                    
-                    with gr.Accordion(label="Advanced Settings", open=False):
-                        with gr.Row(equal_height=True):
-                            guidance_scale = gr.Number(value=1.2, minimum=0.8, maximum=2.5, step=0.1, label="guidance_scale")
-                            steps = gr.Number(value=12, minimum=6, maximum=20, step=1, label="steps")
-                            # strength = gr.Number(value=1, minimum=0.01, maximum=1.0, step=0.01, label="strength")
-                            seed = gr.Number(value=123456, minimum=0, maximum=999999, step=1, label="seed")
-                            negative_prompt = gr.Textbox(label="negative_prompt", value=default_negative_prompt, placeholder=default_negative_prompt, info="what you don't want to see in the image")
-
-                        
-                with gr.Column():
-                    image_out = gr.Image(label="Output", elem_id="output-img", height=400)
-
+        with gr.Column():
+            image = gr.ImageEditor(sources=["upload"], layers=False, transforms=[], brush=gr.Brush(colors=["#000000"], color_mode="fixed"))
+            with gr.Row(elem_id="prompt-container", equal_height=True):
+                btn = gr.Button("Inpaint!", elem_id="run_button")
             
-
-    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, seed], outputs=[image_out], api_name='run')
-    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, seed], outputs=[image_out])
+            with gr.Accordion(label="Advanced Settings", open=False):
+                with gr.Row(equal_height=True):
+                    guidance_scale = gr.Number(value=1.2, minimum=0.0, maximum=2.5, step=0.1, label="guidance_scale")
+                    seed = gr.Number(value=123456, minimum=0, maximum=999999, step=1, label="seed")
+        
+        with gr.Column():
+            image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+
+    # Button click will trigger the inpainting function (no prompt required)
+    btn.click(fn=predict, inputs=[image, guidance_scale, seed], outputs=[image_out], api_name='run')
+    
 
     gr.HTML(
         """