Views
No views yet
[!NOTE] Contains the NF4 checkpoints (transformerandtext_encoder_2) ofblack-forest-labs/FLUX.1-Canny-dev. Please adhere to the original model licensing!
1# !pip install -U controlnet_aux
2from diffusers import DiffusionPipeline, FluxControlPipeline, FluxTransformer2DModel
3import torch
4from transformers import T5EncoderModel
5from controlnet_aux import CannyDetector
6from diffusers.utils import load_image
7import fire
8
9
10def load_pipeline(four_bit=False):
11 orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
12 if four_bit:
13 print("Using four bit.")
14 transformer = FluxTransformer2DModel.from_pretrained(
15 "sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="transformer", torch_dtype=torch.bfloat16
16 )
17 text_encoder_2 = T5EncoderModel.from_pretrained(
18 "sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="text_encoder_2", torch_dtype=torch.bfloat16
19 )
20 pipeline = FluxControlPipeline.from_pipe(
21 orig_pipeline, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16
22 )
23 else:
24 transformer = FluxTransformer2DModel.from_pretrained(
25 "black-forest-labs/FLUX.1-Canny-dev",
26 subfolder="transformer",
27 revision="refs/pr/1",
28 torch_dtype=torch.bfloat16,
29 )
30 pipeline = FluxControlPipeline.from_pipe(orig_pipeline, transformer=transformer, torch_dtype=torch.bfloat16)
31
32 pipeline.enable_model_cpu_offload()
33 return pipeline
34
35def get_canny(control_image):
36 processor = CannyDetector()
37 control_image = processor(
38 control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024
39 )
40 return control_image
41
42def load_conditions():
43 prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
44 control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")
45 control_image = get_canny(control_image)
46 return prompt, control_image
47
48
49def main(four_bit: bool = False):
50 ckpt_id = "sayakpaul/FLUX.1-Canny-dev-nf4"
51 pipe = load_pipeline(four_bit=four_bit)
52 prompt, control_image = load_conditions()
53 image = pipe(
54 prompt=prompt,
55 control_image=control_image,
56 height=1024,
57 width=1024,
58 num_inference_steps=50,
59 guidance_scale=30.0,
60 max_sequence_length=512,
61 generator=torch.Generator("cpu").manual_seed(0),
62 ).images[0]
63 filename = "output_" + ckpt_id.split("/")[-1].replace(".", "_")
64 filename += "_4bit" if four_bit else ""
65 image.save(f"{filename}.png")
66
67
68if __name__ == "__main__":
69 fire.Fire(main)| Original | NF4 |
|---|---|
 |
 |