Update app.py

app.py

@@ -1,52 +1,178 @@
 import gradio as gr
-from PIL import Image
-import requests
-from io import BytesIO
 import torch
-from torchvision import transforms
-from diffusers import AutoencoderKL, LCMScheduler
-from pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline
-from controlnet import ControlNetModel
-
-# Define helper functions
-def download_image(url):
-    response = requests.get(url)
-    return Image.open(BytesIO(response.content)).convert("RGB")
-
-def load_model():
-    # Load model components
-    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)
-    pipe.to('cuda')
-    return pipe
-
-pipe = load_model()
-
-# Define the inpainting function
-def inpaint(image, mask):
-    # Process image and mask
-    image = image.resize((1024, 1024)).convert("RGB")
-    mask = mask.resize((1024, 1024)).convert("L")
-
-    # Transform to tensor
-    image_transform = transforms.ToTensor()
-    image_tensor = image_transform(image).unsqueeze(0).to('cuda')
-    mask_tensor = image_transform(mask).unsqueeze(0).to('cuda')
-    mask_tensor = (mask_tensor > 0.5).float()  # binarize mask
-
-    # Generate image
-    with torch.no_grad():
-        result = pipe(prompt="A park bench", init_image=image_tensor, mask_image=mask_tensor, num_inference_steps=50).images[0]
-
-    return transforms.ToPILImage()(result.squeeze(0))
-
-# Define the interface
-interface = gr.Interface(fn=inpaint,
-                          inputs=[gr.inputs.Image(type="pil", label="Original Image"), gr.inputs.Image(type="pil", label="Mask Image")],
-                          outputs=gr.outputs.Image(type="pil", label="Inpainted Image"),
-                          title="Stable Diffusion XL ControlNet Inpainting",
-                          description="Upload an image and its corresponding mask to inpaint the specified area.")
-
-if __name__ == "__main__":
-    interface.launch()
+import numpy as np
+import diffusers
+import os
+from PIL import Image
+hf_token = os.environ.get("HF_TOKEN")
+from diffusers import StableDiffusionXLInpaintPipeline, DDIMScheduler, UNet2DConditionModel
+
+ratios_map =  {
+    0.5:{"width":704,"height":1408},
+    0.57:{"width":768,"height":1344},
+    0.68:{"width":832,"height":1216},
+    0.72:{"width":832,"height":1152},
+    0.78:{"width":896,"height":1152},
+    0.82:{"width":896,"height":1088},
+    0.88:{"width":960,"height":1088},
+    0.94:{"width":960,"height":1024},
+    1.00:{"width":1024,"height":1024},
+    1.13:{"width":1088,"height":960},
+    1.21:{"width":1088,"height":896},
+    1.29:{"width":1152,"height":896},
+    1.38:{"width":1152,"height":832},
+    1.46:{"width":1216,"height":832},
+    1.67:{"width":1280,"height":768},
+    1.75:{"width":1344,"height":768},
+    2.00:{"width":1408,"height":704}
+}
+ratios = np.array(list(ratios_map.keys()))
+
+def get_size(init_image):
+    w,h=init_image.size
+    curr_ratio = w/h
+    ind = np.argmin(np.abs(curr_ratio-ratios))
+    ratio = ratios[ind]
+    chosen_ratio  = ratios_map[ratio]
+    w,h = chosen_ratio['width'], chosen_ratio['height']
+
+    return w,h
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+unet = UNet2DConditionModel.from_pretrained(
+    "briaai/BRIA-2.2-Inpainting",
+    subfolder="unet",
+    torch_dtype=torch.float16,
+)
+
+scheduler = DDIMScheduler.from_pretrained("briaai/BRIA-2.3", subfolder="scheduler",clip_sample=False)
+
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "briaai/BRIA-2.3",
+    unet=unet,
+    scheduler=scheduler,
+    torch_dtype=torch.float16,
+    force_zeros_for_empty_prompt=False
+)
+
+pipe = pipe.to(device)
+pipe.force_zeros_for_empty_prompt = False
+
+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 read_content(file_path: str) -> str:
+    """read the content of target file
+    """
+    with open(file_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+
+    return content
+
+def predict(dict, prompt="", negative_prompt="", guidance_scale=5, steps=30, strength=1.0):
+    if negative_prompt == "":
+        negative_prompt = None
+
+    
+    init_image = dict["image"].convert("RGB")#.resize((1024, 1024))
+    mask = dict["mask"].convert("RGB")#.resize((1024, 1024))
+
+    w,h = get_size(init_image)
+
+    init_image = init_image.resize((w, h))
+    mask = mask.resize((w, h))
+    
+    # Resize to nearest ratio ?
+    
+    mask = np.array(mask)
+    mask[mask>0]=255
+    mask = Image.fromarray(mask)
+    
+    output = pipe(prompt = prompt,width=w,height=h, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
+    
+    return output.images[0] #, gr.update(visible=True)
+
+
+css = '''
+.gradio-container{max-width: 1100px !important}
+#image_upload{min-height:400px}
+#image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 400px}
+#mask_radio .gr-form{background:transparent; border: none}
+#word_mask{margin-top: .75em !important}
+#word_mask textarea:disabled{opacity: 0.3}
+.footer {margin-bottom: 45px;margin-top: 35px;text-align: center;border-bottom: 1px solid #e5e5e5}
+.footer>p {font-size: .8rem; display: inline-block; padding: 0 10px;transform: translateY(10px);background: white}
+.dark .footer {border-color: #303030}
+.dark .footer>p {background: #0b0f19}
+.acknowledgments h4{margin: 1.25em 0 .25em 0;font-weight: bold;font-size: 115%}
+#image_upload .touch-none{display: flex}
+@keyframes spin {
+    from {
+        transform: rotate(0deg);
+    }
+    to {
+        transform: rotate(360deg);
+    }
+}
+#share-btn-container {padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; max-width: 13rem; margin-left: auto;}
+div#share-btn-container > div {flex-direction: row;background: black;align-items: center}
+#share-btn-container:hover {background-color: #060606}
+#share-btn {all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.5rem !important; padding-bottom: 0.5rem !important;right:0;}
+#share-btn * {all: unset}
+#share-btn-container div:nth-child(-n+2){width: auto !important;min-height: 0px !important;}
+#share-btn-container .wrap {display: none !important}
+#share-btn-container.hidden {display: none!important}
+#prompt input{width: calc(100% - 160px);border-top-right-radius: 0px;border-bottom-right-radius: 0px;}
+#run_button{position:absolute;margin-top: 11px;right: 0;margin-right: 0.8em;border-bottom-left-radius: 0px;
+    border-top-left-radius: 0px;}
+#prompt-container{margin-top:-18px;}
+#prompt-container .form{border-top-left-radius: 0;border-top-right-radius: 0}
+#image_upload{border-bottom-left-radius: 0px;border-bottom-right-radius: 0px}
+'''
+
+image_blocks = gr.Blocks(css=css, elem_id="total-container")
+with image_blocks as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("## BRIA 2.2")
+        gr.HTML('''
+          <p style="margin-bottom: 10px; font-size: 94%">
+            This is a demo for 
+            <a href="https://huggingface.co/briaai/BRIA-2.2" target="_blank">BRIA 2.2 text-to-image </a>. 
+            BRIA 2.2 improve the generation of humans and illustrations compared to BRIA 2.2 while still trained on licensed data, and so provide full legal liability coverage for copyright and privacy infringement.
+          </p>
+        ''')
+    with gr.Row():
+                with gr.Column():
+                    image = 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=5, minimum=1.0, maximum=10.0, step=0.5, label="guidance_scale")
+                            steps = gr.Number(value=30, minimum=20, maximum=50, step=1, label="steps")
+                            strength = gr.Number(value=1, minimum=0.01, maximum=1.0, step=0.01, label="strength")
+                            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)
+
+            
+
+    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength], outputs=[image_out], api_name='run')
+    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength], outputs=[image_out])
+
+    gr.HTML(
+        """
+            <div class="footer">
+                <p>Model by <a href="https://huggingface.co/diffusers" style="text-decoration: underline;" target="_blank">Diffusers</a> - Gradio Demo by 🤗 Hugging Face
+                </p>
+            </div>
+        """
+    )
+
+image_blocks.queue(max_size=25,api_open=False).launch(show_api=False)