sayakpaul HF staff commited on
Commit
3f69266
1 Parent(s): 4827d45

Create README.md

Browse files
Files changed (1) hide show
  1. README.md +107 -0
README.md ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: diffusers
3
+ license: other
4
+ license_name: flux-1-dev-non-commercial-license
5
+ license_link: LICENSE.md
6
+ ---
7
+
8
+ > [!NOTE]
9
+ > Contains the NF4 checkpoints (`transformer` and `text_encoder_2`) of [`black-forest-labs/FLUX.1-Canny-dev`](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev). Please adhere to the original model licensing!
10
+
11
+ <details>
12
+ <summary>Code</summary>
13
+
14
+ ```py
15
+ # !pip install -U controlnet_aux
16
+ from diffusers import DiffusionPipeline, FluxControlPipeline, FluxTransformer2DModel
17
+ import torch
18
+ from transformers import T5EncoderModel
19
+ from controlnet_aux import CannyDetector
20
+ from diffusers.utils import load_image
21
+ import fire
22
+
23
+
24
+ def load_pipeline(four_bit=False):
25
+ orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
26
+ if four_bit:
27
+ print("Using four bit.")
28
+ transformer = FluxTransformer2DModel.from_pretrained(
29
+ "sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="transformer", torch_dtype=torch.bfloat16
30
+ )
31
+ text_encoder_2 = T5EncoderModel.from_pretrained(
32
+ "sayakpaul/FLUX.1-Canny-dev-nf4", subfolder="text_encoder_2", torch_dtype=torch.bfloat16
33
+ )
34
+ pipeline = FluxControlPipeline.from_pipe(
35
+ orig_pipeline, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16
36
+ )
37
+ else:
38
+ transformer = FluxTransformer2DModel.from_pretrained(
39
+ "black-forest-labs/FLUX.1-Canny-dev",
40
+ subfolder="transformer",
41
+ revision="refs/pr/1",
42
+ torch_dtype=torch.bfloat16,
43
+ )
44
+ pipeline = FluxControlPipeline.from_pipe(orig_pipeline, transformer=transformer, torch_dtype=torch.bfloat16)
45
+
46
+ pipeline.enable_model_cpu_offload()
47
+ return pipeline
48
+
49
+ def get_canny(control_image):
50
+ processor = CannyDetector()
51
+ control_image = processor(
52
+ control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024
53
+ )
54
+ return control_image
55
+
56
+ def load_conditions():
57
+ prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
58
+ control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")
59
+ control_image = get_canny(control_image)
60
+ return prompt, control_image
61
+
62
+
63
+ def main(four_bit: bool = False):
64
+ ckpt_id = "sayakpaul/FLUX.1-Canny-dev-nf4"
65
+ pipe = load_pipeline(four_bit=four_bit)
66
+ prompt, control_image = load_conditions()
67
+ image = pipe(
68
+ prompt=prompt,
69
+ control_image=control_image,
70
+ height=1024,
71
+ width=1024,
72
+ num_inference_steps=50,
73
+ guidance_scale=30.0,
74
+ max_sequence_length=512,
75
+ generator=torch.Generator("cpu").manual_seed(0),
76
+ ).images[0]
77
+ filename = "output_" + ckpt_id.split("/")[-1].replace(".", "_")
78
+ filename += "_4bit" if four_bit else ""
79
+ image.save(f"{filename}.png")
80
+
81
+
82
+ if __name__ == "__main__":
83
+ fire.Fire(main)
84
+ ```
85
+
86
+ </details>
87
+
88
+ ## Outputs
89
+
90
+ <table>
91
+ <thead>
92
+ <tr>
93
+ <th>Original</th>
94
+ <th>NF4</th>
95
+ </tr>
96
+ </thead>
97
+ <tbody>
98
+ <tr>
99
+ <td>
100
+ <img src="./assets/output_FLUX_1-Canny-dev.png" alt="Original">
101
+ </td>
102
+ <td>
103
+ <img src="./assets/output_FLUX_1-Canny-dev_4bit.png" alt="NF4">
104
+ </td>
105
+ </tr>
106
+ </tbody>
107
+ </table>