Update app.py

app.py

@@ -38,6 +38,27 @@ model_configs = {
     'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
 }
 
+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()))
+
 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")
@@ -45,27 +66,59 @@ state_dict = torch.load(filepath, map_location="cpu")
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()
 
+from huggingface_hub import hf_hub_download
+import os
+
+try:
+    local_dir = os.path.dirname(__file__)
+except:
+    local_dir = '.'
+    
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='pipeline_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='transformer_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-4B-Adapt", filename='bria_utils.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.0-ControlNet-Union", filename='pipeline_bria_controlnet.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.0-ControlNet-Union", filename='controlnet_bria.py', local_dir=local_dir)
+
+
 import torch
 from diffusers.utils import load_image
-from diffusers import FluxControlNetPipeline, FluxControlNetModel
-from diffusers.models import FluxMultiControlNetModel
-
-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)
+from controlnet_bria import BriaControlNetModel, BriaMultiControlNetModel
+from pipeline_bria_controlnet import BriaControlNetPipeline
+import PIL.Image as Image
+
+base_model = 'briaai/BRIA-4B-Adapt'
+controlnet_model = 'briaai/BRIA-3.0-ControlNet-Union'
+
+controlnet = BriaControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+controlnet = BriaMultiControlNetModel([controlnet])
+
+pipe = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
 pipe.to("cuda")
 
-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}
+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,
+}
 
 canny = CannyDetector()
 open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
 
 torch.backends.cuda.matmul.allow_tf32 = True
-pipe.vae.enable_tiling()
-pipe.vae.enable_slicing()
+# 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:
@@ -100,13 +153,8 @@ def convert_to_grayscale(image):
     gray_image = convert_from_cv2_to_image(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
     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):
+def tile_old(input_image, resolution=768):
     input_image = convert_from_image_to_cv2(input_image)
     H, W, C = input_image.shape
     H = float(H)
@@ -114,38 +162,32 @@ def tile(input_image, resolution=768):
     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
+    H = int(np.round(H / 16.0)) * 16
+    W = int(np.round(W / 16.0)) * 16
     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):
+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)
+    
+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
 
-    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 resize_image(image):
+    image = image.convert('RGB')
+    w,h = get_size(image)
+    resized_image = image.resize((w, h))
+    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)):
-        
     control_mode_num = mode_mapping[control_mode]
     
     if cond_in is None:
@@ -157,14 +199,12 @@ def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mo
                 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":
+            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(image_in)
+                control_image = tile(16, image_in)
     else:
         control_image = resize_img(load_image(cond_in))
 
@@ -214,7 +254,7 @@ with gr.Blocks(css=css) as demo:
                     
                     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="gray",
                             info="select the control mode, one for all"
                         )
                         
@@ -223,7 +263,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(