Update app.py

app.py

@@ -5,58 +5,12 @@ 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']
+from diffusers import AutoPipelineForText2Image
 
-    return w,h
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-unet = UNet2DConditionModel.from_pretrained(
-    "briaai/BRIA-2.3-Inpainting",
-    subfolder="unet",
-    torch_dtype=torch.float16,
-)
-
-scheduler = DDIMScheduler.from_pretrained("briaai/BRIA-2.3", subfolder="scheduler", 
-                                                        rescale_betas_zero_snr=True,prediction_type='v_prediction',timestep_spacing="trailing",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
+pipe = AutoPipelineForText2Image.from_pretrained("briaai/BRIA-2.3", torch_dtype=torch.float16, force_zeros_for_empty_prompt=False).to(device)
+pipe.load_ip_adapter("briaai/DEV-Image-Prompt", subfolder='models', weight_name="ip_adapter_bria.bin")
 
 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"
 
@@ -69,27 +23,28 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(dict, prompt="", negative_prompt="", guidance_scale=5, steps=30, strength=1.0):
-    if negative_prompt == "":
-        negative_prompt = None
+def predict(dict, prompt="high quality, best quality", negative_prompt="", guidance_scale=5, steps=30, ip_adapter_scale = 1.0, width=1024, height=1024, seed=0):
 
-    
-    init_image = dict["image"].convert("RGB")#.resize((1024, 1024))
-    mask = dict["mask"].convert("RGB")#.resize((1024, 1024))
-
-    w,h = get_size(init_image)
+    pipeline.set_ip_adapter_scale(ip_adapter_scale)
 
-    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)
+    if negative_prompt == "":
+        negative_prompt = None
     
-    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)
+    init_image = dict["image"].convert("RGB")
+    init_image = init_image.resize((224, 224))
     
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+        
+    output = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    ip_adapter_image=init_image,
+    num_inference_steps=steps,
+    generator=generator,
+    height=height, width=width,
+    guidance_scale=guidance_scale
+    ).images
+
     return output.images[0] #, gr.update(visible=True)
 
 
@@ -146,14 +101,20 @@ with image_blocks as demo:
                     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")
+                            prompt = gr.Textbox(placeholder="Your prompt (you can leave it empty if you only want the image prompt as input)", show_label=False, elem_id="prompt")
+                            btn = gr.Button("Generate!", elem_id="run_button")
+
+                    with gr.Accordion(label="Settings", open=True):
+                        with gr.Row(equal_height=True):
+                            ip_adapter_scale = gr.Number(value=1.0, minimum=0.01, maximum=1.0, step=0.01, label="ip_adapter_scale")
+                            width = gr.Number(value=1024, minimum=0.01, maximum=1.0, step=0.01, label="width")
+                            height = gr.Number(value=1024, minimum=0.01, maximum=1.0, step=0.01, label="height")                    
                     
                     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")
+                            steps = gr.Number(value=30, minimum=10, maximum=100, step=1, label="steps")
+                            seed = gr.Number(value=0, minimum=0, maximum=100000, 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")
 
                         
@@ -162,8 +123,8 @@ with image_blocks as demo:
 
             
 
-    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])
+    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, ip_adapter_scale, width, height, seed], outputs=[image_out], api_name='run')
+    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, ip_adapter_scale, width, height, seed], outputs=[image_out])
 
     # gr.Examples(
     #             examples=[