Update app.py

app.py

@@ -14,17 +14,28 @@ import random
 
 os.system("pip install -e ./controlnet_aux")
 
-from controlnet_aux import OpenposeDetector, CannyDetector
+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_GATED")
+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.2", filename='pipeline_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2", filename='transformer_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2", filename='bria_utils.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2-ControlNet-Union", filename='pipeline_bria_controlnet.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2-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)
@@ -38,6 +49,20 @@ model_configs = {
     '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")
@@ -47,25 +72,35 @@ model = model.to(DEVICE).eval()
 
 import torch
 from diffusers.utils import load_image
-from diffusers import FluxControlNetPipeline, FluxControlNetModel
-from diffusers.models import FluxMultiControlNetModel
+from controlnet_bria import BriaControlNetModel, BriaMultiControlNetModel
+from pipeline_bria_controlnet import BriaControlNetPipeline
+import PIL.Image as Image
 
-base_model = 'black-forest-labs/FLUX.1-dev'
-controlnet_model = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro'
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
-controlnet = FluxMultiControlNetModel([controlnet])
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
-pipe.to("cuda")
+base_model = 'briaai/BRIA-3.2'
+controlnet_model = 'briaai/BRIA-3.2-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 = {"canny":0, "tile":1, "depth":2, "blur":3, "openpose":4, "gray":5, "low quality": 6}
-strength_mapping = {"canny":0.65, "tile":0.45, "depth":0.55, "blur":0.45, "openpose":0.55, "gray":0.45, "low quality": 0.4}
-
-canny = CannyDetector()
+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.vae.enable_tiling()
-pipe.vae.enable_slicing()
 pipe.enable_model_cpu_offload() # for saving memory
 
 def convert_from_image_to_cv2(img: Image) -> np.ndarray:
@@ -84,102 +119,69 @@ def extract_depth(image):
 
 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(image):
-    processed_image_canny = canny(image)
-    return processed_image_canny
+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 apply_gaussian_blur(image, kernel_size=(21, 21)):
-    image = convert_from_image_to_cv2(image)
-    blurred_image = convert_from_cv2_to_image(cv2.GaussianBlur(image, kernel_size, 0))
-    return blurred_image
 
 def convert_to_grayscale(image):
-    image = convert_from_image_to_cv2(image)
-    gray_image = convert_from_cv2_to_image(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
+    gray_image = image.convert('L').convert('RGB')
     return gray_image
 
-def add_gaussian_noise(image, mean=0, sigma=10):
-    image = convert_from_image_to_cv2(image)
-    noise = np.random.normal(mean, sigma, image.shape)
-    noisy_image = convert_from_cv2_to_image(np.clip(image.astype(np.float32) + noise, 0, 255).astype(np.uint8))
-    return noisy_image
-
-def tile(input_image, resolution=768):
-    input_image = convert_from_image_to_cv2(input_image)
-    H, W, C = input_image.shape
-    H = float(H)
-    W = float(W)
-    k = float(resolution) / min(H, W)
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
-    img = convert_from_cv2_to_image(img)
-    return img
-
-def resize_img(input_image, max_side=768, min_side=512, 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 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(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
-        
+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 cond_in is None:
-        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 == "openpose":
-                control_image = extract_openpose(image_in)
-            elif control_mode == "blur":
-                control_image = apply_gaussian_blur(image_in)
-            elif control_mode == "low quality":
-                control_image = add_gaussian_noise(image_in)
-            elif control_mode == "gray":
-                control_image = convert_to_grayscale(image_in)
-            elif control_mode == "tile":
-                control_image = tile(image_in)
-    else:
-        control_image = resize_img(load_image(cond_in))
+    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],
+        control_image=control_image,
+        control_mode=control_mode_num,
         width=width,
         height=height,
-        controlnet_conditioning_scale=[control_strength],
+        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() 
@@ -196,9 +198,8 @@ css="""
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
-        # FLUX.1-dev-ControlNet-Union-Pro
-        A unified ControlNet for FLUX.1-dev model from the InstantX team and Shakker Labs. Model card: [Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro](https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro). <br />
-        The recommended strength: {"canny":0.65, "tile":0.45, "depth":0.55, "blur":0.45, "openpose":0.55, "gray":0.45, "low quality": 0.4}. Long prompt is preferred by FLUX.1.
+        # BRIA-3.2-ControlNet-Union
+        A unified ControlNet for BRIA-3.2 model from Bria.ai.<br />
         """)
         
         with gr.Column():
@@ -206,15 +207,15 @@ with gr.Blocks(css=css) as demo:
             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")
+                    # 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(
-                            ["canny", "depth", "openpose", "gray", "blur", "tile", "low quality"], label="Mode", value="gray",
+                            ["depth", "canny", "colorgrid", "recolor", "tile", "pose"], label="Mode", value="canny",
                             info="select the control mode, one for all"
                         )
                         
@@ -223,7 +224,7 @@ with gr.Blocks(css=css) as demo:
                             minimum=0,
                             maximum=1.0,
                             step=0.05,
-                            value=0.50,
+                            value=0.9,
                         )
                     
                     seed = gr.Slider(
@@ -231,15 +232,15 @@ with gr.Blocks(css=css) as demo:
                         minimum=0,
                         maximum=MAX_SEED,
                         step=1,
-                        value=42,
+                        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=24)
-                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
+                                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")
                     
@@ -255,7 +256,7 @@ with gr.Blocks(css=css) as demo:
         api_name=False
     ).then(
         fn = infer,
-        inputs = [cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
+        inputs = [image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
         outputs = [result, processed_cond],
         show_api=False
     )