Update README.md

README.md

@@ -118,38 +118,62 @@ image = pipe(
 
 # Multi-Controls Inference
 ```python
+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, 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_union = 'InstantX/FLUX.1-dev-Controlnet-Union'
+base_model = 'briaai/BRIA-4B-Adapt'
+controlnet_model = 'briaai/BRIA-3.0-ControlNet-Union'
 
-controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch.bfloat16)
-controlnet = FluxMultiControlNetModel([controlnet_union]) # we always recommend loading via FluxMultiControlNetModel
+controlnet = BriaControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+controlnet = BriaMultiControlNetModel([controlnet])
 
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
+pipe = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
 pipe.to("cuda")
 
-prompt = 'A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.'
-control_image_depth = load_image("https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/resolve/main/images/depth.jpg")
-control_mode_depth = 2
+# control_image_colorgrid = load_image("https://huggingface.co/briaai/BRIA-3.0-ControlNet-Union/resolve/main/colorgrid.jpg")
+# control_image_pose = load_image("https://huggingface.co/briaai/BRIA-3.0-ControlNet-Union/resolve/main/pose.jpg")
+
+control_image_colorgrid = Image.open("/home/ubuntu/spring/Infra/sd3/eval_results/control_images/colorgrid/38.jpg").convert("RGB")
+control_image_pose = Image.open("/home/ubuntu/spring/Infra/sd3/eval_results/control_images/pose/38.jpg").convert("RGB")
+
+control_image = [control_image_colorgrid, control_image_pose]
+controlnet_conditioning_scale = [0.5, 0.5]
+control_mode = [2, 5]
 
-control_image_canny = load_image("https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/resolve/main/images/canny.jpg")
-control_mode_canny = 0
+width, height = control_image[0].size
 
-width, height = control_image.size
+prompt = 'Two kids in jackets play near a tent in a forest with trees.'
 
+generator = torch.Generator(device="cuda").manual_seed(555)
 image = pipe(
     prompt, 
-    control_image=[control_image_depth, control_image_canny],
-    control_mode=[control_mode_depth, control_mode_canny],
+    control_image=control_image,
+    control_mode=control_mode,
     width=width,
     height=height,
-    controlnet_conditioning_scale=[0.2, 0.4],
-    num_inference_steps=24, 
-    guidance_scale=3.5,
-    generator=torch.manual_seed(42),
+    controlnet_conditioning_scale=controlnet_conditioning_scale,
+    num_inference_steps=50, 
+    max_sequence_length=128,
+    guidance_scale=5,
+    generator=generator
 ).images[0]
 
 ```