README.md

---
license: other
license_name: flux-non-commercial-license
base_model: black-forest-labs/FLUX.2-klein-base-9B
library_name: diffusers
pipeline_tag: image-to-image
tags:
- flux
- semantic-segmentation
- instance-segmentation
- referring-segmentation
- depth-estimation
- surface-normal-estimation
---

## Quickstart

Install PyTorch with CUDA support first, then:

```shell
pip install "diffusers>=0.37.0" transformers accelerate safetensors pillow
```

The fine-tuned transformer is swapped into the [FLUX.2 [klein] base 9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) pipeline. One model handles every task and the prompt selects it. This example runs semantic segmentation:

```python
import torch
from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel
from diffusers.utils import load_image
from PIL import Image

device = "cuda"
dtype = torch.bfloat16

transformer = Flux2Transformer2DModel.from_pretrained(
    "FunAILab/Open-Vision-Banana",
    subfolder="final_models/9B/nonzeroshot/checkpoint-avg82500-85000-87500-92500-97500_uniform/transformer",
    torch_dtype=dtype,
)
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-base-9B", transformer=transformer, torch_dtype=dtype
)
pipe.enable_model_cpu_offload()  # save some VRAM by offloading the model to CPU

image = load_image("street.jpg")
prompt = (
    "[TASK: SEMANTIC_SEGMENTATION] Generate a semantic segmentation visualization of "
    'the input image, using this color mapping: {"road": [8,4,249], "sidewalk": [255,8,232], '
    '"vegetation": [10,255,2], "sky": [2,204,254], "person": [255,0,13], "car": [255,249,23]} '
    "Color each pixel of a listed class with its color and leave everything else black. "
    "Output only the segmentation visualization image."
)

# The pipeline runs at the input's aspect ratio, capped at 1 megapixel.
output = pipe(
    prompt=prompt,
    image=image,
    guidance_scale=3.5,
    num_inference_steps=28,
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]
# Back to the input size: NEAREST keeps mask colours exact; use Image.BILINEAR for depth and normals.
output.resize(image.size, Image.NEAREST).save("output.png")
```

### Quickstart prompts for the other tasks

Replace `prompt` above with one of these.

**Referring segmentation**

```python
prompt = (
    "[TASK: REFERRING_SEGMENTATION] Segment the region that best matches this "
    'description: "table in the room". Paint the matching region red [220, 30, 30] and '
    "leave everything else black [0, 0, 0]. Output only the segmentation mask image."
)
```

**Instance segmentation**

```python
prompt = (
    "[TASK: INSTANCE_SEGMENTATION] This image shows segmentation masks for car from the input image. "
    "The background is set to black color [0, 0, 0]. Each car instance is represented by a solid mask, "
    "and a different color is used for each mask. Instances are numbered by where their mask first "
    "appears in a left-to-right scan of the image (ties broken top to bottom). The 1st car is blue, "
    "the 2nd is lime, the 3rd is red, the 4th is sky blue, the 5th is dark green, the 6th is pink, "
    "the 7th is apricot, the 8th is mint, the 9th is wine, the 10th is magenta, the 11th is navy, "
    "the 12th is green, the 13th is yellow, the 14th is brown, the 15th is white, the 16th is olive, "
    "the 17th is lavender, the 18th is teal, the 19th is orange, the 20th is hot pink, the 21st is plum, "
    "the 22nd is gray, the 23rd is purple, the 24th is beige, the 25th is steel blue, the 26th is cyan, "
    "the 27th is royal blue, the 28th is sea green, the 29th is rust, the 30th is mustard, "
    "the 31st is mauve, the 32nd is maroon."
)
```

**Depth estimation**

```python
prompt = (
    "Estimate a dense relative inverse-depth map for the input image. Produce exactly one output value "
    "for every input pixel at the same pixel location, with closer points brighter and farther points "
    "darker. Encode the result as neutral grayscale (R=G=B) using one consistent mapping over the "
    "complete image. Preserve object silhouettes, thin structures, occlusion boundaries, smooth surfaces, "
    "and small depth changes. Do not alter the scene or add texture, shading, color, text, labels, "
    "legends, borders, or objects; return only the registered depth map."
)
```

**Surface normal estimation**

```python
prompt = (
    "Create a dense camera-space unit surface-normal map registered to the input. Use +x image-right, "
    "+y image-up, +z toward the camera and encode (x,y,z) as R=(1-x)/2, G=(1+y)/2, B=(1+z)/2. "
    "Preserve planes, curved surfaces, thin structures, and sharp orientation boundaries; "
    "return only the RGB normal map."
)
```

These are quickstart prompts, one per task. The model was trained on many phrasings of each task, so you can adapt them:

- **Semantic segmentation:** list any classes with an RGB colour each. Distinct, saturated colours work best. Pixels outside the listed classes come out black.
- **Referring segmentation:** replace the description and the paint colour.
- **Instance segmentation:** prompt one class at a time; replace every `car` in the prompt with any noun phrase. The left-to-right color ordering helps maintain greater separation between assigned colors. The model may not follow it exactly, which is fine: each instance only needs its own distinct colour.
- **Depth and surface normals:** use the prompts as they are. Depth is relative inverse depth, with closer points brighter. Normals follow the encoding stated in the prompt.