Update README.md

README.md

@@ -48,7 +48,7 @@ pip install diffusers==0.30.2, hf_hub_download
 ```python
 from huggingface_hub import hf_hub_download
 import os
-
+import pandas as pd
 try:
     local_dir = os.path.dirname(__file__)
 except:
@@ -62,27 +62,49 @@ hf_hub_download(repo_id="briaai/BRIA-4B-Adapt-ControlNet-Union", filename='contr
 
 import torch
 from diffusers.utils import load_image
-from controlnet_bria import BriaControlNetModel, BriaMultiControlNetModel
+from controlnet_bria import BriaControlNetModel
 from pipeline_bria_controlnet import BriaControlNetPipeline
 import PIL.Image as Image
 
+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},
+}
+
+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
+
+
 base_model = 'briaai/BRIA-4B-Adapt'
 controlnet_model = 'briaai/BRIA-4B-Adapt-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.to("cuda")
+pipeline = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
+pipeline = pipeline.to(device="cuda", dtype=torch.bfloat16)
 
 control_image_canny = load_image("https://huggingface.co/briaai/BRIA-4B-Adapt-ControlNet-Union/resolve/main/images/canny.jpg")
 controlnet_conditioning_scale = 1.0
 control_mode = 1
-
+control_image_canny = resize_img(control_image_canny)
 width, height = control_image_canny.size
 
 prompt = 'In a serene living room, someone rests on a sapphire blue couch, diligently drawing in a rose-tinted notebook, with a sleek black coffee table, a muted green wall, an elegant geometric lamp, and a lush potted palm enhancing the peaceful ambiance.'
 
 generator = torch.Generator(device="cuda").manual_seed(555)
-image = pipe(
+image = pipeline(
     prompt, 
     control_image=control_image_canny,
     control_mode=control_mode,
@@ -92,8 +114,10 @@ image = pipe(
     num_inference_steps=50, 
     max_sequence_length=128,
     guidance_scale=5,
-    generator=generator
+    generator=generator,
+    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]
+print(image)        
 ```
 
 # Multi-Controls Inference