Create README.md

README.md

@@ -0,0 +1,106 @@
+---
+library_name: diffusers
+license: other
+license_name: flux-1-dev-non-commercial-license
+license_link: LICENSE.md
+---
+
+> [!NOTE]
+> Contains the NF4 checkpoints (`transformer` and `text_encoder_2`) of [`black-forest-labs/FLUX.1-Depth-dev`](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev). Please adhere to the original model licensing!
+
+<details>
+  <summary>Code</summary>
+
+```py
+# !pip install git+https://github.com/asomoza/image_gen_aux.git
+from diffusers import DiffusionPipeline, FluxControlPipeline, FluxTransformer2DModel
+import torch
+from transformers import T5EncoderModel
+from image_gen_aux import DepthPreprocessor
+from diffusers.utils import load_image
+import fire
+
+
+def load_pipeline(four_bit=False):
+    orig_pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+    if four_bit:
+        print("Using four bit.")
+        transformer = FluxTransformer2DModel.from_pretrained(
+            "sayakpaul/FLUX.1-Depth-dev-nf4", subfolder="transformer", torch_dtype=torch.bfloat16
+        )
+        text_encoder_2 = T5EncoderModel.from_pretrained(
+            "sayakpaul/FLUX.1-Depth-dev-nf4", subfolder="text_encoder_2", torch_dtype=torch.bfloat16
+        )
+        pipeline = FluxControlPipeline.from_pipe(
+            orig_pipeline, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16
+        )
+    else:
+        transformer = FluxTransformer2DModel.from_pretrained(
+            "black-forest-labs/FLUX.1-Depth-dev",
+            subfolder="transformer",
+            revision="refs/pr/1",
+            torch_dtype=torch.bfloat16,
+        )
+        pipeline = FluxControlPipeline.from_pipe(orig_pipeline, transformer=transformer, torch_dtype=torch.bfloat16)
+
+    pipeline.enable_model_cpu_offload()
+    return pipeline
+
+@torch.no_grad()
+def get_depth(control_image):
+    processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
+    control_image = processor(control_image)[0].convert("RGB")
+    return control_image
+
+def load_conditions():
+    prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
+    control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")
+    control_image = get_depth(control_image)
+    return prompt, control_image
+
+
+def main(four_bit: bool = False):
+    ckpt_id = "sayakpaul/FLUX.1-Depth-dev-nf4"
+    pipe = load_pipeline(ckpt_id=ckpt_id, four_bit=four_bit)
+    prompt, control_image = load_conditions()
+    image = pipe(
+        prompt=prompt,
+        control_image=control_image,
+        height=1024,
+        width=1024,
+        num_inference_steps=30,
+        guidance_scale=10.0,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0),
+    ).images[0]
+    filename = "output_" + ckpt_id.split("/")[-1].replace(".", "_")
+    filename += "_4bit" if four_bit else ""
+    image.save(f"{filename}.png")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)
+```
+
+</details>
+
+## Outputs
+
+<table>
+    <thead>
+        <tr>
+            <th>Original</th>
+            <th>NF4</th>
+        </tr>
+    </thead>
+    <tbody>
+        <tr>
+            <td>
+                <img src="./assets/output_FLUX_1-Fill-dev.png" alt="Original">
+            </td>
+            <td>
+                <img src="./assets/output_FLUX_1-Fill-dev_4bit.png" alt="NF4">
+            </td>
+        </tr>
+    </tbody>
+</table>