Create app.py

app.py

@@ -0,0 +1,457 @@
+import os
+import random
+import gradio as gr
+
+
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
+from transformers import CLIPVisionModelWithProjection
+from diffusers.utils import load_image
+from diffusers.models import ControlNetModel
+# from diffusers.image_processor import IPAdapterMaskProcessor
+from insightface.app import FaceAnalysis
+# import sys
+# import glob
+# import os
+import io
+import spaces
+
+from pipeline_stable_diffusion_xl_instantid import StableDiffusionXLInstantIDPipeline, draw_kps
+
+import pandas as pd
+import json
+import requests
+from PIL import Image
+from io import BytesIO
+
+
+def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
+               pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):
+    
+    w, h = input_image.size
+    if size is not None:
+        w_resize_new, h_resize_new = size
+    else:
+        ratio = min_side / min(h, w)
+        w, h = round(ratio*w), round(ratio*h)
+        ratio = max_side / max(h, w)
+        input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
+        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
+        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
+    input_image = input_image.resize([w_resize_new, h_resize_new], mode)
+
+    if pad_to_max_side:
+        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
+        offset_x = (max_side - w_resize_new) // 2
+        offset_y = (max_side - h_resize_new) // 2
+        res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
+        input_image = Image.fromarray(res)
+    return input_image
+
+def process_image_by_bbox_larger(input_image, bbox_xyxy, min_bbox_ratio=0.2):
+    """
+    Process an image based on a bounding box, cropping and resizing as necessary.
+
+    Parameters:
+    - input_image: PIL Image object.
+    - bbox_xyxy: Tuple (x1, y1, x2, y2) representing the bounding box coordinates.
+
+    Returns:
+    - A processed image cropped and resized to 1024x1024 if the bounding box is valid,
+      or None if the bounding box does not meet the required size criteria.
+    """
+    # Constants
+    target_size = 1024
+    # min_bbox_ratio = 0.2  # Bounding box should be at least 20% of the crop
+
+    # Extract bounding box coordinates
+    x1, y1, x2, y2 = bbox_xyxy
+    bbox_w = x2 - x1
+    bbox_h = y2 - y1
+
+    # Calculate the area of the bounding box
+    bbox_area = bbox_w * bbox_h
+
+    # Start with the smallest square crop that allows bbox to be at least 20% of the crop area
+    crop_size = max(bbox_w, bbox_h)
+    initial_crop_area = crop_size * crop_size
+    while (bbox_area / initial_crop_area) < min_bbox_ratio:
+        crop_size += 10  # Gradually increase until bbox is at least 20% of the area
+        initial_crop_area = crop_size * crop_size
+
+    # Once the minimum condition is satisfied, try to expand the crop further
+    max_possible_crop_size = min(input_image.width, input_image.height)
+    while crop_size < max_possible_crop_size:
+        # Calculate a potential new area
+        new_crop_size = crop_size + 10
+        new_crop_area = new_crop_size * new_crop_size
+        if (bbox_area / new_crop_area) < min_bbox_ratio:
+            break  # Stop if expanding further violates the 20% rule
+        crop_size = new_crop_size
+
+    # Determine the center of the bounding box
+    center_x = (x1 + x2) // 2
+    center_y = (y1 + y2) // 2
+
+    # Calculate the crop coordinates centered around the bounding box
+    crop_x1 = max(0, center_x - crop_size // 2)
+    crop_y1 = max(0, center_y - crop_size // 2)
+    crop_x2 = min(input_image.width, crop_x1 + crop_size)
+    crop_y2 = min(input_image.height, crop_y1 + crop_size)
+
+    # Ensure the crop is square, adjust if it goes out of image bounds
+    if crop_x2 - crop_x1 != crop_y2 - crop_y1:
+        side_length = min(crop_x2 - crop_x1, crop_y2 - crop_y1)
+        crop_x2 = crop_x1 + side_length
+        crop_y2 = crop_y1 + side_length
+
+    # Crop the image
+    cropped_image = input_image.crop((crop_x1, crop_y1, crop_x2, crop_y2))
+
+    # Resize the cropped image to 1024x1024
+    resized_image = cropped_image.resize((target_size, target_size), Image.LANCZOS)
+
+    return resized_image
+
+def calc_emb_cropped(image, app):
+    face_image = image.copy()
+
+    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
+
+    face_info = face_info[0]
+
+    cropped_face_image = process_image_by_bbox_larger(face_image, face_info["bbox"], min_bbox_ratio=0.2)
+
+    return cropped_face_image
+
+def process_benchmark_csv(banchmark_csv_path):
+    # Reading the first CSV file into a DataFrame
+    df = pd.read_csv(banchmark_csv_path)
+
+    # Drop any unnamed columns
+    df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
+
+    # Drop columns with all NaN values
+    df.dropna(axis=1, how='all', inplace=True)
+
+    # Drop rows with all NaN values
+    df.dropna(axis=0, how='all', inplace=True)
+
+    df = df.loc[df['High resolution'] == 1]
+
+    df.reset_index(drop=True, inplace=True)
+
+    return df
+
+def make_canny_condition(image, min_val=100, max_val=200, w_bilateral=True):
+    if w_bilateral:
+        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+        bilateral_filtered_image = cv2.bilateralFilter(image, d=9, sigmaColor=75, sigmaSpace=75)
+        image = cv2.Canny(bilateral_filtered_image, min_val, max_val)
+    else:
+        image = np.array(image)
+        image = cv2.Canny(image, min_val, max_val)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    image = Image.fromarray(image)
+    return image
+
+
+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"
+
+# Load face detection and recognition package
+app = FaceAnalysis(name='antelopev2', root='./', providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+base_dir = "./instantID_ckpt/checkpoint_174000"
+face_adapter = f'{base_dir}/pytorch_model.bin'
+controlnet_path = f'{base_dir}/controlnet'
+base_model_path = f'briaai/BRIA-2.3'
+resolution = 1024
+
+controlnet_lnmks = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
+
+controlnet_canny = ControlNetModel.from_pretrained("briaai/BRIA-2.3-ControlNet-Canny",
+                                                   torch_dtype=torch.float16)
+
+controlnet = [controlnet_lnmks, controlnet_canny]
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+        '/home/ubuntu/BRIA-2.3-InstantID/ip_adapter/image_encoder',
+        torch_dtype=torch.float16,
+    )
+
+pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
+        base_model_path,
+        controlnet=controlnet,
+        torch_dtype=torch.float16,
+        image_encoder=image_encoder # For compatibility issues - needs to be there
+    )
+
+pipe = pipe.to(device)
+
+use_native_ip_adapter = True
+pipe.use_native_ip_adapter=use_native_ip_adapter
+
+pipe.load_ip_adapter_instantid(face_adapter)
+
+clip_embeds=None
+
+
+Loras_dict = {
+    "":"",
+    "Vangogh_Vanilla": "bold, dramatic brush strokes, vibrant colors, swirling patterns, intense, emotionally charged paintings of",
+    "Avatar_internlm": "2d anime sketch avatar of",
+    # "Tomer_Hanuka_V3": "Fluid lines",
+    "Storyboards": "Illustration style for storyboarding",
+    "3D_illustration": "3D object illustration, abstract",
+    # "beetl_general_death_style_v2": "a pale, dead, unnatural color face with dark circles around the eyes",
+    "Characters": "gaming vector Art"
+    }
+
+lora_names = Loras_dict.keys()
+
+lora_base_path = "./LoRAs"
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, 99999999)
+    return seed
+
+
+@spaces.GPU
+def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale=0.8, kps_scale=0.6, canny_scale=0.4, lora_name="", lora_scale=0.7, progress=gr.Progress(track_tqdm=True)):
+    if image_path is None:
+        raise gr.Error(f"Cannot find any input face image! Please upload a face image.")
+    
+    # img = np.array(Image.open(image_path))[:,:,::-1]
+    img = Image.open(image_path)
+
+    face_image_orig = img #Image.open(BytesIO(response.content))
+    face_image_cropped = calc_emb_cropped(face_image_orig, app)
+    face_image = resize_img(face_image_cropped, max_side=resolution, min_side=resolution)
+    # face_image_padded = resize_img(face_image_cropped, max_side=resolution, min_side=resolution, pad_to_max_side=True)
+    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
+    face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
+    face_emb = face_info['embedding']
+    face_kps = draw_kps(face_image, face_info['kps'])
+
+    if canny_scale>0.0:
+        # Convert PIL image to a file-like object
+        image_file = io.BytesIO()
+        face_image_cropped.save(image_file, format='JPEG')  # Save in the desired format (e.g., 'JPEG' or 'PNG')
+        image_file.seek(0)  # Move to the start of the BytesIO stream
+
+        url = "https://engine.prod.bria-api.com/v1/background/remove"
+
+        payload = {}
+        files = [
+            ('file', ('image_name.jpeg', image_file, 'image/jpeg'))  # Specify file name, file-like object, and MIME type
+        ]
+        headers = {
+        'api_token': 'a10d6386dd6a11ebba800242ac130004'
+        }
+
+        response = requests.request("POST", url, headers=headers, data=payload, files=files)
+
+        print(response.text)
+
+        response_json = json.loads(response.content.decode('utf-8'))
+
+        img = requests.get(response_json['result_url'])
+
+        processed_image = Image.open(io.BytesIO(img.content))
+
+        # Assuming `processed_image` is the RGBA image returned
+        if processed_image.mode == 'RGBA':
+            # Create a white background image
+            white_background = Image.new("RGB", processed_image.size, (255, 255, 255))
+            # Composite the RGBA image over the white background
+            face_image = Image.alpha_composite(white_background.convert('RGBA'), processed_image).convert('RGB')
+        else:
+            face_image = processed_image.convert('RGB')  # If already RGB, just ensure mode is correct
+
+        canny_img = make_canny_condition(face_image, min_val=20, max_val=40, w_bilateral=True)
+                    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    if lora_name != "":
+        lora_path = os.path.join(lora_base_path, lora_name, "pytorch_lora_weights.safetensors")
+        pipe.load_lora_weights(lora_path)
+        pipe.fuse_lora(lora_scale)  
+        pipe.enable_lora()
+
+        lora_prefix = Loras_dict[lora_name]
+
+        prompt = f"{lora_prefix} {prompt}"
+        
+
+    print("Start inference...")    
+    images = pipe(
+        prompt = prompt,
+        negative_prompt = default_negative_prompt,
+        image_embeds = face_emb,
+        image = [face_kps, canny_img] if canny_scale>0.0 else face_kps,
+        controlnet_conditioning_scale = [kps_scale, canny_scale] if canny_scale>0.0 else kps_scale,
+        control_guidance_end = [1.0, 1.0] if canny_scale>0.0 else 1.0,
+        ip_adapter_scale = ip_adapter_scale,
+        num_inference_steps = num_steps,
+        guidance_scale = guidance_scale,
+        generator = generator,
+        visual_prompt_embds = clip_embeds,
+        cross_attention_kwargs = None,
+        num_images_per_prompt=num_images,
+    ).images #[0]
+
+    if lora_name != "":
+        pipe.disable_lora()
+        pipe.unfuse_lora()
+        pipe.unload_lora_weights()
+
+    return images
+
+### Description
+title = r"""
+<h1>Bria-2.3 ID preservation</h1>
+"""
+
+description = r"""
+<b>🤗 Gradio demo</b> for bria ID preservation.<br>
+
+Steps:<br>
+1. Upload an image with a face. If multiple faces are detected, we use the largest one. For images with already tightly cropped faces, detection may fail, try images with a larger margin.
+2. Click <b>Submit</b> to generate new images of the subject.
+"""
+
+Footer = r"""
+Enjoy
+"""
+
+css = '''
+.gradio-container {width: 85% !important}
+'''
+with gr.Blocks(css=css) as demo:
+
+    # description
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Column():
+            
+            # upload face image
+            img_file = gr.Image(label="Upload a photo with a face", type="filepath")
+            
+            # Textbox for entering a prompt
+            prompt = gr.Textbox(
+            label="Prompt", 
+            placeholder="Enter your prompt here", 
+            info="Describe what you want to generate or modify in the image."
+            )
+            
+            lora_name = gr.Dropdown(choices=lora_names, label="LoRA", value="",  info="Select a LoRA name from the list, not selecting any will disable LoRA.")
+
+            submit = gr.Button("Submit", variant="primary")
+
+            # use_lcm = gr.Checkbox(
+            #     label="Use LCM-LoRA to accelerate sampling", value=False,
+            #     info="Reduces sampling steps significantly, but may decrease quality.",
+            # )
+            
+            with gr.Accordion(open=False, label="Advanced Options"):
+                num_steps = gr.Slider( 
+                    label="Number of sample steps",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=30,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.1,
+                    maximum=10.0,
+                    step=0.1,
+                    value=5.0,
+                )
+                num_images = gr.Slider(
+                    label="Number of output images",
+                    minimum=1,
+                    maximum=3,
+                    step=1,
+                    value=1,
+                )
+                ip_adapter_scale = gr.Slider(
+                    label="ip adapter scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.8,
+                )
+                kps_scale = gr.Slider(
+                    label="kps control scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.6,
+                )
+                canny_scale = gr.Slider(
+                    label="canny control scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.4,
+                )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=99999999,
+                    step=1,
+                    value=0,
+                )
+                seed = gr.Slider(
+                    label="lora_scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.7,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        with gr.Column():
+            gallery = gr.Gallery(label="Generated Images")
+
+        submit.click(
+            fn=randomize_seed_fn,
+            inputs=[seed, randomize_seed],
+            outputs=seed,
+            queue=False,
+            api_name=False,
+        ).then(
+            fn=generate_image,
+            inputs=[img_file, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name],
+            outputs=[gallery]
+        )
+
+    # use_lcm.input(
+    #         fn=toggle_lcm_ui,
+    #         inputs=[use_lcm],
+    #         outputs=[num_steps, guidance_scale],
+    #         queue=False,
+    #     )       
+    
+    # gr.Examples(
+    #     examples=get_example(),
+    #     inputs=[img_file],
+    #     run_on_click=True,
+    #     fn=run_example,
+    #     outputs=[gallery],
+    # )
+    
+    gr.Markdown(Footer)
+
+demo.launch()