Create app.py

app.py

@@ -0,0 +1,265 @@
+import sys
+sys.path.append('./')
+
+import gradio as gr
+import spaces
+import os
+import sys
+import subprocess
+import numpy as np
+from PIL import Image
+import cv2
+import torch
+import random
+
+os.system("pip install -e ./controlnet_aux")
+
+from controlnet_aux import OpenposeDetector #, CannyDetector
+from depth_anything_v2.dpt import DepthAnythingV2
+
+from huggingface_hub import hf_hub_download
+
+from huggingface_hub import login
+hf_token = os.environ.get("HF_TOKEN")
+login(token=hf_token)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+try:
+    local_dir = os.path.dirname(__file__)
+except:
+    local_dir = '.'
+    
+hf_hub_download(repo_id="briaai/BRIA-3.1", filename='pipeline_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.1", filename='transformer_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.1", filename='bria_utils.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.1-ControlNet-Union", filename='pipeline_bria_controlnet.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.1-ControlNet-Union", filename='controlnet_bria.py', local_dir=local_dir)
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+model_configs = {
+    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
+    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
+}
+
+RATIO_CONFIGS_1024 = {
+    0.6666666666666666: {"width": 832, "height": 1248},
+    0.7432432432432432: {"width": 880, "height": 1184},
+    0.8028169014084507: {"width": 912, "height": 1136},
+    1.0: {"width": 1024, "height": 1024},
+    1.2456140350877194: {"width": 1136, "height": 912},
+    1.3454545454545455: {"width": 1184, "height": 880},
+    1.4339622641509433: {"width": 1216, "height": 848},
+    1.5: {"width": 1248, "height": 832},
+    1.5490196078431373: {"width": 1264, "height": 816},
+    1.62: {"width": 1296, "height": 800},
+    1.7708333333333333: {"width": 1360, "height": 768},
+}
+
+encoder = 'vitl'
+model = DepthAnythingV2(**model_configs[encoder])
+filepath = hf_hub_download(repo_id=f"depth-anything/Depth-Anything-V2-Large", filename=f"depth_anything_v2_vitl.pth", repo_type="model")
+state_dict = torch.load(filepath, map_location="cpu")
+model.load_state_dict(state_dict)
+model = model.to(DEVICE).eval()
+
+import torch
+from diffusers.utils import load_image
+from controlnet_bria import BriaControlNetModel, BriaMultiControlNetModel
+from pipeline_bria_controlnet import BriaControlNetPipeline
+import PIL.Image as Image
+
+base_model = 'briaai/BRIA-3.1'
+controlnet_model = 'briaai/BRIA-3.1-ControlNet-Union'
+controlnet = BriaControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+pipe = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
+pipe = pipe.to(device="cuda", dtype=torch.bfloat16)
+
+mode_mapping = {
+    "depth": 0,
+    "canny": 1,
+    "colorgrid": 2,
+    "recolor": 3,
+    "tile": 4,
+    "pose": 5,
+}
+strength_mapping = {
+    "depth": 1.0,
+    "canny": 1.0,
+    "colorgrid": 1.0,
+    "recolor": 1.0,
+    "tile": 1.0,
+    "pose": 1.0,
+}
+open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+
+torch.backends.cuda.matmul.allow_tf32 = True
+pipe.enable_model_cpu_offload() # for saving memory
+
+def convert_from_image_to_cv2(img: Image) -> np.ndarray:
+    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+def convert_from_cv2_to_image(img: np.ndarray) -> Image:
+    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+def extract_depth(image):
+    image = np.asarray(image)
+    depth = model.infer_image(image[:, :, ::-1])
+    depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
+    depth = depth.astype(np.uint8)
+    gray_depth = Image.fromarray(depth).convert('RGB') 
+    return gray_depth
+
+def extract_openpose(img):
+    processed_image_open_pose = open_pose(img, hand_and_face=True)
+    processed_image_open_pose = processed_image_open_pose.resize(img.size)
+    return processed_image_open_pose
+    
+def extract_canny(input_image):
+    image = np.array(input_image)
+    image = cv2.Canny(image, 100, 200)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    canny_image = Image.fromarray(image)
+    return canny_image
+
+
+def convert_to_grayscale(image):
+    gray_image = image.convert('L').convert('RGB')
+    return gray_image
+
+def tile(downscale_factor, input_image):
+    control_image = input_image.resize((input_image.size[0] // downscale_factor, input_image.size[1] // downscale_factor)).resize(input_image.size, Image.NEAREST)
+    return control_image
+    
+def resize_img(control_image):
+    image_ratio = control_image.width / control_image.height
+    ratio = min(RATIO_CONFIGS_1024.keys(), key=lambda k: abs(k - image_ratio))
+    to_height = RATIO_CONFIGS_1024[ratio]["height"]
+    to_width = RATIO_CONFIGS_1024[ratio]["width"]
+    resized_image = control_image.resize((to_width, to_height), resample=Image.Resampling.LANCZOS)
+    return resized_image
+
+@spaces.GPU(duration=180)
+def infer(image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
+    control_mode_num = mode_mapping[control_mode]
+    
+    if image_in is not None:
+        image_in = resize_img(load_image(image_in))
+        if control_mode == "canny":
+            control_image = extract_canny(image_in)
+        elif control_mode == "depth":
+            control_image = extract_depth(image_in)
+        elif control_mode == "pose":
+            control_image = extract_openpose(image_in)
+        elif control_mode == "colorgrid":
+            control_image = tile(64, image_in)
+        elif control_mode == "recolor":
+            control_image = convert_to_grayscale(image_in)
+        elif control_mode == "tile":
+            control_image = tile(16, image_in)
+
+    control_image = resize_img(control_image)
+
+    width, height = control_image.size
+    
+    image = pipe(
+        prompt, 
+        control_image=control_image,
+        control_mode=control_mode_num,
+        width=width,
+        height=height,
+        controlnet_conditioning_scale=control_strength,
+        num_inference_steps=inference_steps, 
+        guidance_scale=guidance_scale,
+        generator=torch.manual_seed(seed),
+        max_sequence_length=128,
+        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"
+    ).images[0]
+
+    torch.cuda.empty_cache() 
+    
+    return image, control_image, gr.update(visible=True)
+   
+
+css="""
+#col-container{
+    margin: 0 auto;
+    max-width: 1080px;
+}
+"""
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("""
+        # BRIA-3.1-ControlNet-Union
+        A unified ControlNet for BRIA-3.1 model from Bria.ai.<br />
+        """)
+        
+        with gr.Column():
+            
+            with gr.Row():
+                with gr.Column():
+                    
+                    # with gr.Row(equal_height=True):
+                        # cond_in = gr.Image(label="Upload a processed control image", sources=["upload"], type="filepath")
+                    image_in = gr.Image(label="Extract condition from a reference image (Optional)", sources=["upload"], type="filepath")
+                    
+                    prompt = gr.Textbox(label="Prompt", value="best quality")
+                    
+                    with gr.Accordion("Controlnet"):
+                        control_mode = gr.Radio(
+                            ["depth", "canny", "colorgrid", "recolor", "tile", "pose"], label="Mode", value="canny",
+                            info="select the control mode, one for all"
+                        )
+                        
+                        control_strength = gr.Slider(
+                            label="control strength",
+                            minimum=0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=0.9,
+                        )
+                    
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=555,
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                    
+                    with gr.Accordion("Advanced settings", open=False):
+                        with gr.Column():
+                            with gr.Row():
+                                inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=50)
+                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0)
+                    
+                    submit_btn = gr.Button("Submit")
+                    
+                with gr.Column():
+                    result = gr.Image(label="Result")
+                    processed_cond = gr.Image(label="Preprocessed Cond")
+
+    submit_btn.click(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False
+    ).then(
+        fn = infer,
+        inputs = [image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
+        outputs = [result, processed_cond],
+        show_api=False
+    )
+
+demo.queue(api_open=False)
+demo.launch()