Delete flux

flux/__init__.py

@@ -1,13 +0,0 @@
-try:
-    from ._version import (
-        version as __version__,  # type: ignore
-        version_tuple,
-    )
-except ImportError:
-    __version__ = "unknown (no version information available)"
-    version_tuple = (0, 0, "unknown", "noinfo")
-
-from pathlib import Path
-
-PACKAGE = __package__.replace("_", "-")
-PACKAGE_ROOT = Path(__file__).parent

flux/__main__.py

@@ -1,4 +0,0 @@
-from .cli import app
-
-if __name__ == "__main__":
-    app()

flux/api.py

@@ -1,225 +0,0 @@
-import io
-import os
-import time
-from pathlib import Path
-
-import requests
-from PIL import Image
-
-API_URL = "https://api.bfl.ml"
-API_ENDPOINTS = {
-    "flux.1-pro": "flux-pro",
-    "flux.1-dev": "flux-dev",
-    "flux.1.1-pro": "flux-pro-1.1",
-}
-
-
-class ApiException(Exception):
-    def __init__(self, status_code: int, detail: str | list[dict] | None = None):
-        super().__init__()
-        self.detail = detail
-        self.status_code = status_code
-
-    def __str__(self) -> str:
-        return self.__repr__()
-
-    def __repr__(self) -> str:
-        if self.detail is None:
-            message = None
-        elif isinstance(self.detail, str):
-            message = self.detail
-        else:
-            message = "[" + ",".join(d["msg"] for d in self.detail) + "]"
-        return f"ApiException({self.status_code=}, {message=}, detail={self.detail})"
-
-
-class ImageRequest:
-    def __init__(
-        self,
-        # api inputs
-        prompt: str,
-        name: str = "flux.1.1-pro",
-        width: int | None = None,
-        height: int | None = None,
-        num_steps: int | None = None,
-        prompt_upsampling: bool | None = None,
-        seed: int | None = None,
-        guidance: float | None = None,
-        interval: float | None = None,
-        safety_tolerance: int | None = None,
-        # behavior of this class
-        validate: bool = True,
-        launch: bool = True,
-        api_key: str | None = None,
-    ):
-        """
-        Manages an image generation request to the API.
-
-        All parameters not specified will use the API defaults.
-
-        Args:
-            prompt: Text prompt for image generation.
-            width: Width of the generated image in pixels. Must be a multiple of 32.
-            height: Height of the generated image in pixels. Must be a multiple of 32.
-            name: Which model version to use
-            num_steps: Number of steps for the image generation process.
-            prompt_upsampling: Whether to perform upsampling on the prompt.
-            seed: Optional seed for reproducibility.
-            guidance: Guidance scale for image generation.
-            safety_tolerance: Tolerance level for input and output moderation.
-                 Between 0 and 6, 0 being most strict, 6 being least strict.
-            validate: Run input validation
-            launch: Directly launches request
-            api_key: Your API key if not provided by the environment
-
-        Raises:
-            ValueError: For invalid input, when `validate`
-            ApiException: For errors raised from the API
-        """
-        if validate:
-            if name not in API_ENDPOINTS.keys():
-                raise ValueError(f"Invalid model {name}")
-            elif width is not None and width % 32 != 0:
-                raise ValueError(f"width must be divisible by 32, got {width}")
-            elif width is not None and not (256 <= width <= 1440):
-                raise ValueError(f"width must be between 256 and 1440, got {width}")
-            elif height is not None and height % 32 != 0:
-                raise ValueError(f"height must be divisible by 32, got {height}")
-            elif height is not None and not (256 <= height <= 1440):
-                raise ValueError(f"height must be between 256 and 1440, got {height}")
-            elif num_steps is not None and not (1 <= num_steps <= 50):
-                raise ValueError(f"steps must be between 1 and 50, got {num_steps}")
-            elif guidance is not None and not (1.5 <= guidance <= 5.0):
-                raise ValueError(f"guidance must be between 1.5 and 4, got {guidance}")
-            elif interval is not None and not (1.0 <= interval <= 4.0):
-                raise ValueError(f"interval must be between 1 and 4, got {interval}")
-            elif safety_tolerance is not None and not (0 <= safety_tolerance <= 6.0):
-                raise ValueError(f"safety_tolerance must be between 0 and 6, got {interval}")
-
-            if name == "flux.1-dev":
-                if interval is not None:
-                    raise ValueError("Interval is not supported for flux.1-dev")
-            if name == "flux.1.1-pro":
-                if interval is not None or num_steps is not None or guidance is not None:
-                    raise ValueError("Interval, num_steps and guidance are not supported for " "flux.1.1-pro")
-
-        self.name = name
-        self.request_json = {
-            "prompt": prompt,
-            "width": width,
-            "height": height,
-            "steps": num_steps,
-            "prompt_upsampling": prompt_upsampling,
-            "seed": seed,
-            "guidance": guidance,
-            "interval": interval,
-            "safety_tolerance": safety_tolerance,
-        }
-        self.request_json = {key: value for key, value in self.request_json.items() if value is not None}
-
-        self.request_id: str | None = None
-        self.result: dict | None = None
-        self._image_bytes: bytes | None = None
-        self._url: str | None = None
-        if api_key is None:
-            self.api_key = os.environ.get("BFL_API_KEY")
-        else:
-            self.api_key = api_key
-
-        if launch:
-            self.request()
-
-    def request(self):
-        """
-        Request to generate the image.
-        """
-        if self.request_id is not None:
-            return
-        response = requests.post(
-            f"{API_URL}/v1/{API_ENDPOINTS[self.name]}",
-            headers={
-                "accept": "application/json",
-                "x-key": self.api_key,
-                "Content-Type": "application/json",
-            },
-            json=self.request_json,
-        )
-        result = response.json()
-        if response.status_code != 200:
-            raise ApiException(status_code=response.status_code, detail=result.get("detail"))
-        self.request_id = response.json()["id"]
-
-    def retrieve(self) -> dict:
-        """
-        Wait for the generation to finish and retrieve response.
-        """
-        if self.request_id is None:
-            self.request()
-        while self.result is None:
-            response = requests.get(
-                f"{API_URL}/v1/get_result",
-                headers={
-                    "accept": "application/json",
-                    "x-key": self.api_key,
-                },
-                params={
-                    "id": self.request_id,
-                },
-            )
-            result = response.json()
-            if "status" not in result:
-                raise ApiException(status_code=response.status_code, detail=result.get("detail"))
-            elif result["status"] == "Ready":
-                self.result = result["result"]
-            elif result["status"] == "Pending":
-                time.sleep(0.5)
-            else:
-                raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'")
-        return self.result
-
-    @property
-    def bytes(self) -> bytes:
-        """
-        Generated image as bytes.
-        """
-        if self._image_bytes is None:
-            response = requests.get(self.url)
-            if response.status_code == 200:
-                self._image_bytes = response.content
-            else:
-                raise ApiException(status_code=response.status_code)
-        return self._image_bytes
-
-    @property
-    def url(self) -> str:
-        """
-        Public url to retrieve the image from
-        """
-        if self._url is None:
-            result = self.retrieve()
-            self._url = result["sample"]
-        return self._url
-
-    @property
-    def image(self) -> Image.Image:
-        """
-        Load the image as a PIL Image
-        """
-        return Image.open(io.BytesIO(self.bytes))
-
-    def save(self, path: str):
-        """
-        Save the generated image to a local path
-        """
-        suffix = Path(self.url).suffix
-        if not path.endswith(suffix):
-            path = path + suffix
-        Path(path).resolve().parent.mkdir(parents=True, exist_ok=True)
-        with open(path, "wb") as file:
-            file.write(self.bytes)
-
-
-if __name__ == "__main__":
-    from fire import Fire
-
-    Fire(ImageRequest)

flux/cli.py

@@ -1,238 +0,0 @@
-import os
-import re
-import time
-from dataclasses import dataclass
-from glob import iglob
-
-import torch
-from fire import Fire
-from transformers import pipeline
-
-from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
-from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
-
-NSFW_THRESHOLD = 0.85
-
-
-@dataclass
-class SamplingOptions:
-    prompt: str
-    width: int
-    height: int
-    num_steps: int
-    guidance: float
-    seed: int | None
-
-
-def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
-    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the prompt or write a command starting with a slash:\n"
-        "- '/w <width>' will set the width of the generated image\n"
-        "- '/h <height>' will set the height of the generated image\n"
-        "- '/s <seed>' sets the next seed\n"
-        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
-        "- '/n <steps>' sets the number of steps\n"
-        "- '/q' to quit"
-    )
-
-    while (prompt := input(user_question)).startswith("/"):
-        if prompt.startswith("/w"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, width = prompt.split()
-            options.width = 16 * (int(width) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/h"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, height = prompt.split()
-            options.height = 16 * (int(height) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/g"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, guidance = prompt.split()
-            options.guidance = float(guidance)
-            print(f"Setting guidance to {options.guidance}")
-        elif prompt.startswith("/s"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, seed = prompt.split()
-            options.seed = int(seed)
-            print(f"Setting seed to {options.seed}")
-        elif prompt.startswith("/n"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, steps = prompt.split()
-            options.num_steps = int(steps)
-            print(f"Setting number of steps to {options.num_steps}")
-        elif prompt.startswith("/q"):
-            print("Quitting")
-            return None
-        else:
-            if not prompt.startswith("/h"):
-                print(f"Got invalid command '{prompt}'\n{usage}")
-            print(usage)
-    if prompt != "":
-        options.prompt = prompt
-    return options
-
-
-@torch.inference_mode()
-def main(
-    name: str = "flux-schnell",
-    width: int = 1360,
-    height: int = 768,
-    seed: int | None = None,
-    prompt: str = (
-        "a photo of a forest with mist swirling around the tree trunks. The word "
-        '"FLUX" is painted over it in big, red brush strokes with visible texture'
-    ),
-    device: str = "cuda" if torch.cuda.is_available() else "cpu",
-    num_steps: int | None = None,
-    loop: bool = False,
-    guidance: float = 3.5,
-    offload: bool = False,
-    output_dir: str = "output",
-    add_sampling_metadata: bool = True,
-):
-    """
-    Sample the flux model. Either interactively (set `--loop`) or run for a
-    single image.
-
-    Args:
-        name: Name of the model to load
-        height: height of the sample in pixels (should be a multiple of 16)
-        width: width of the sample in pixels (should be a multiple of 16)
-        seed: Set a seed for sampling
-        output_name: where to save the output image, `{idx}` will be replaced
-            by the index of the sample
-        prompt: Prompt used for sampling
-        device: Pytorch device
-        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
-        loop: start an interactive session and sample multiple times
-        guidance: guidance value used for guidance distillation
-        add_sampling_metadata: Add the prompt to the image Exif metadata
-    """
-    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
-
-    if name not in configs:
-        available = ", ".join(configs.keys())
-        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
-
-    torch_device = torch.device(device)
-    if num_steps is None:
-        num_steps = 4 if name == "flux-schnell" else 50
-
-    # allow for packing and conversion to latent space
-    height = 16 * (height // 16)
-    width = 16 * (width // 16)
-
-    output_name = os.path.join(output_dir, "img_{idx}.jpg")
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        idx = 0
-    else:
-        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
-        if len(fns) > 0:
-            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
-        else:
-            idx = 0
-
-    # init all components
-    t5 = load_t5(torch_device, max_length=256 if name == "flux-schnell" else 512)
-    clip = load_clip(torch_device)
-    model = load_flow_model(name, device="cpu" if offload else torch_device)
-    ae = load_ae(name, device="cpu" if offload else torch_device)
-
-    rng = torch.Generator(device="cpu")
-    opts = SamplingOptions(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_steps=num_steps,
-        guidance=guidance,
-        seed=seed,
-    )
-
-    if loop:
-        opts = parse_prompt(opts)
-
-    while opts is not None:
-        if opts.seed is None:
-            opts.seed = rng.seed()
-        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
-        t0 = time.perf_counter()
-
-        # prepare input
-        x = get_noise(
-            1,
-            opts.height,
-            opts.width,
-            device=torch_device,
-            dtype=torch.bfloat16,
-            seed=opts.seed,
-        )
-        opts.seed = None
-        if offload:
-            ae = ae.cpu()
-            torch.cuda.empty_cache()
-            t5, clip = t5.to(torch_device), clip.to(torch_device)
-        inp = prepare(t5, clip, x, prompt=opts.prompt)
-        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
-
-        # offload TEs to CPU, load model to gpu
-        if offload:
-            t5, clip = t5.cpu(), clip.cpu()
-            torch.cuda.empty_cache()
-            model = model.to(torch_device)
-
-        # denoise initial noise
-        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
-
-        # offload model, load autoencoder to gpu
-        if offload:
-            model.cpu()
-            torch.cuda.empty_cache()
-            ae.decoder.to(x.device)
-
-        # decode latents to pixel space
-        x = unpack(x.float(), opts.height, opts.width)
-        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
-            x = ae.decode(x)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        t1 = time.perf_counter()
-
-        fn = output_name.format(idx=idx)
-        print(f"Done in {t1 - t0:.1f}s. Saving {fn}")
-
-        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
-
-        if loop:
-            print("-" * 80)
-            opts = parse_prompt(opts)
-        else:
-            opts = None
-
-
-def app():
-    Fire(main)
-
-
-if __name__ == "__main__":
-    app()

flux/cli_control.py

@@ -1,347 +0,0 @@
-import os
-import re
-import time
-from dataclasses import dataclass
-from glob import iglob
-
-import torch
-from fire import Fire
-from transformers import pipeline
-
-from flux.modules.image_embedders import CannyImageEncoder, DepthImageEncoder
-from flux.sampling import denoise, get_noise, get_schedule, prepare_control, unpack
-from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
-
-
-@dataclass
-class SamplingOptions:
-    prompt: str
-    width: int
-    height: int
-    num_steps: int
-    guidance: float
-    seed: int | None
-    img_cond_path: str
-    lora_scale: float | None
-
-
-def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
-    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the prompt or write a command starting with a slash:\n"
-        "- '/w <width>' will set the width of the generated image\n"
-        "- '/h <height>' will set the height of the generated image\n"
-        "- '/s <seed>' sets the next seed\n"
-        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
-        "- '/n <steps>' sets the number of steps\n"
-        "- '/q' to quit"
-    )
-
-    while (prompt := input(user_question)).startswith("/"):
-        if prompt.startswith("/w"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, width = prompt.split()
-            options.width = 16 * (int(width) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/h"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, height = prompt.split()
-            options.height = 16 * (int(height) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/g"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, guidance = prompt.split()
-            options.guidance = float(guidance)
-            print(f"Setting guidance to {options.guidance}")
-        elif prompt.startswith("/s"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, seed = prompt.split()
-            options.seed = int(seed)
-            print(f"Setting seed to {options.seed}")
-        elif prompt.startswith("/n"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, steps = prompt.split()
-            options.num_steps = int(steps)
-            print(f"Setting number of steps to {options.num_steps}")
-        elif prompt.startswith("/q"):
-            print("Quitting")
-            return None
-        else:
-            if not prompt.startswith("/h"):
-                print(f"Got invalid command '{prompt}'\n{usage}")
-            print(usage)
-    if prompt != "":
-        options.prompt = prompt
-    return options
-
-
-def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
-    if options is None:
-        return None
-
-    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the conditioning image or write a command starting with a slash:\n"
-        "- '/q' to quit"
-    )
-
-    while True:
-        img_cond_path = input(user_question)
-
-        if img_cond_path.startswith("/"):
-            if img_cond_path.startswith("/q"):
-                print("Quitting")
-                return None
-            else:
-                if not img_cond_path.startswith("/h"):
-                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
-                print(usage)
-            continue
-
-        if img_cond_path == "":
-            break
-
-        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
-            (".jpg", ".jpeg", ".png", ".webp")
-        ):
-            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
-            continue
-
-        options.img_cond_path = img_cond_path
-        break
-
-    return options
-
-
-def parse_lora_scale(options: SamplingOptions | None) -> tuple[SamplingOptions | None, bool]:
-    changed = False
-
-    if options is None:
-        return None, changed
-
-    user_question = "Next lora scale (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the lora scale or write a command starting with a slash:\n"
-        "- '/q' to quit"
-    )
-
-    while (prompt := input(user_question)).startswith("/"):
-        if prompt.startswith("/q"):
-            print("Quitting")
-            return None, changed
-        else:
-            if not prompt.startswith("/h"):
-                print(f"Got invalid command '{prompt}'\n{usage}")
-            print(usage)
-    if prompt != "":
-        options.lora_scale = float(prompt)
-        changed = True
-    return options, changed
-
-
-@torch.inference_mode()
-def main(
-    name: str,
-    width: int = 1024,
-    height: int = 1024,
-    seed: int | None = None,
-    prompt: str = "a robot made out of gold",
-    device: str = "cuda" if torch.cuda.is_available() else "cpu",
-    num_steps: int = 50,
-    loop: bool = False,
-    guidance: float | None = None,
-    offload: bool = False,
-    output_dir: str = "output",
-    add_sampling_metadata: bool = True,
-    img_cond_path: str = "assets/robot.webp",
-    lora_scale: float | None = 0.85,
-):
-    """
-    Sample the flux model. Either interactively (set `--loop`) or run for a
-    single image.
-
-    Args:
-        height: height of the sample in pixels (should be a multiple of 16)
-        width: width of the sample in pixels (should be a multiple of 16)
-        seed: Set a seed for sampling
-        output_name: where to save the output image, `{idx}` will be replaced
-            by the index of the sample
-        prompt: Prompt used for sampling
-        device: Pytorch device
-        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
-        loop: start an interactive session and sample multiple times
-        guidance: guidance value used for guidance distillation
-        add_sampling_metadata: Add the prompt to the image Exif metadata
-        img_cond_path: path to conditioning image (jpeg/png/webp)
-    """
-    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
-
-    assert name in [
-        "flux-dev-canny",
-        "flux-dev-depth",
-        "flux-dev-canny-lora",
-        "flux-dev-depth-lora",
-    ], f"Got unknown model name: {name}"
-    if guidance is None:
-        if name in ["flux-dev-canny", "flux-dev-canny-lora"]:
-            guidance = 30.0
-        elif name in ["flux-dev-depth", "flux-dev-depth-lora"]:
-            guidance = 10.0
-        else:
-            raise NotImplementedError()
-
-    if name not in configs:
-        available = ", ".join(configs.keys())
-        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
-
-    torch_device = torch.device(device)
-
-    output_name = os.path.join(output_dir, "img_{idx}.jpg")
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        idx = 0
-    else:
-        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
-        if len(fns) > 0:
-            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
-        else:
-            idx = 0
-
-    # init all components
-    t5 = load_t5(torch_device, max_length=512)
-    clip = load_clip(torch_device)
-    model = load_flow_model(name, device="cpu" if offload else torch_device)
-    ae = load_ae(name, device="cpu" if offload else torch_device)
-
-    # set lora scale
-    if "lora" in name and lora_scale is not None:
-        for _, module in model.named_modules():
-            if hasattr(module, "set_scale"):
-                module.set_scale(lora_scale)
-
-    if name in ["flux-dev-depth", "flux-dev-depth-lora"]:
-        img_embedder = DepthImageEncoder(torch_device)
-    elif name in ["flux-dev-canny", "flux-dev-canny-lora"]:
-        img_embedder = CannyImageEncoder(torch_device)
-    else:
-        raise NotImplementedError()
-
-    rng = torch.Generator(device="cpu")
-    opts = SamplingOptions(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_steps=num_steps,
-        guidance=guidance,
-        seed=seed,
-        img_cond_path=img_cond_path,
-        lora_scale=lora_scale,
-    )
-
-    if loop:
-        opts = parse_prompt(opts)
-        opts = parse_img_cond_path(opts)
-        if "lora" in name:
-            opts, changed = parse_lora_scale(opts)
-            if changed:
-                # update the lora scale:
-                for _, module in model.named_modules():
-                    if hasattr(module, "set_scale"):
-                        module.set_scale(opts.lora_scale)
-
-    while opts is not None:
-        if opts.seed is None:
-            opts.seed = rng.seed()
-        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
-        t0 = time.perf_counter()
-
-        # prepare input
-        x = get_noise(
-            1,
-            opts.height,
-            opts.width,
-            device=torch_device,
-            dtype=torch.bfloat16,
-            seed=opts.seed,
-        )
-        opts.seed = None
-        if offload:
-            t5, clip, ae = t5.to(torch_device), clip.to(torch_device), ae.to(torch_device)
-        inp = prepare_control(
-            t5,
-            clip,
-            x,
-            prompt=opts.prompt,
-            ae=ae,
-            encoder=img_embedder,
-            img_cond_path=opts.img_cond_path,
-        )
-        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
-
-        # offload TEs and AE to CPU, load model to gpu
-        if offload:
-            t5, clip, ae = t5.cpu(), clip.cpu(), ae.cpu()
-            torch.cuda.empty_cache()
-            model = model.to(torch_device)
-
-        # denoise initial noise
-        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
-
-        # offload model, load autoencoder to gpu
-        if offload:
-            model.cpu()
-            torch.cuda.empty_cache()
-            ae.decoder.to(x.device)
-
-        # decode latents to pixel space
-        x = unpack(x.float(), opts.height, opts.width)
-        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
-            x = ae.decode(x)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        t1 = time.perf_counter()
-        print(f"Done in {t1 - t0:.1f}s")
-
-        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
-
-        if loop:
-            print("-" * 80)
-            opts = parse_prompt(opts)
-            opts = parse_img_cond_path(opts)
-            if "lora" in name:
-                opts, changed = parse_lora_scale(opts)
-                if changed:
-                    # update the lora scale:
-                    for _, module in model.named_modules():
-                        if hasattr(module, "set_scale"):
-                            module.set_scale(opts.lora_scale)
-        else:
-            opts = None
-
-
-def app():
-    Fire(main)
-
-
-if __name__ == "__main__":
-    app()

flux/cli_fill.py

@@ -1,334 +0,0 @@
-import os
-import re
-import time
-from dataclasses import dataclass
-from glob import iglob
-
-import torch
-from fire import Fire
-from PIL import Image
-from transformers import pipeline
-
-from flux.sampling import denoise, get_noise, get_schedule, prepare_fill, unpack
-from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
-
-
-@dataclass
-class SamplingOptions:
-    prompt: str
-    width: int
-    height: int
-    num_steps: int
-    guidance: float
-    seed: int | None
-    img_cond_path: str
-    img_mask_path: str
-
-
-def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
-    user_question = "Next prompt (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the prompt or write a command starting with a slash:\n"
-        "- '/s <seed>' sets the next seed\n"
-        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
-        "- '/n <steps>' sets the number of steps\n"
-        "- '/q' to quit"
-    )
-
-    while (prompt := input(user_question)).startswith("/"):
-        if prompt.startswith("/g"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, guidance = prompt.split()
-            options.guidance = float(guidance)
-            print(f"Setting guidance to {options.guidance}")
-        elif prompt.startswith("/s"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, seed = prompt.split()
-            options.seed = int(seed)
-            print(f"Setting seed to {options.seed}")
-        elif prompt.startswith("/n"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, steps = prompt.split()
-            options.num_steps = int(steps)
-            print(f"Setting number of steps to {options.num_steps}")
-        elif prompt.startswith("/q"):
-            print("Quitting")
-            return None
-        else:
-            if not prompt.startswith("/h"):
-                print(f"Got invalid command '{prompt}'\n{usage}")
-            print(usage)
-    if prompt != "":
-        options.prompt = prompt
-    return options
-
-
-def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
-    if options is None:
-        return None
-
-    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the conditioning image or write a command starting with a slash:\n"
-        "- '/q' to quit"
-    )
-
-    while True:
-        img_cond_path = input(user_question)
-
-        if img_cond_path.startswith("/"):
-            if img_cond_path.startswith("/q"):
-                print("Quitting")
-                return None
-            else:
-                if not img_cond_path.startswith("/h"):
-                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
-                print(usage)
-            continue
-
-        if img_cond_path == "":
-            break
-
-        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
-            (".jpg", ".jpeg", ".png", ".webp")
-        ):
-            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
-            continue
-        else:
-            with Image.open(img_cond_path) as img:
-                width, height = img.size
-
-            if width % 32 != 0 or height % 32 != 0:
-                print(f"Image dimensions must be divisible by 32, got {width}x{height}")
-                continue
-
-        options.img_cond_path = img_cond_path
-        break
-
-    return options
-
-
-def parse_img_mask_path(options: SamplingOptions | None) -> SamplingOptions | None:
-    if options is None:
-        return None
-
-    user_question = "Next conditioning mask (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the conditioning mask or write a command starting with a slash:\n"
-        "- '/q' to quit"
-    )
-
-    while True:
-        img_mask_path = input(user_question)
-
-        if img_mask_path.startswith("/"):
-            if img_mask_path.startswith("/q"):
-                print("Quitting")
-                return None
-            else:
-                if not img_mask_path.startswith("/h"):
-                    print(f"Got invalid command '{img_mask_path}'\n{usage}")
-                print(usage)
-            continue
-
-        if img_mask_path == "":
-            break
-
-        if not os.path.isfile(img_mask_path) or not img_mask_path.lower().endswith(
-            (".jpg", ".jpeg", ".png", ".webp")
-        ):
-            print(f"File '{img_mask_path}' does not exist or is not a valid image file")
-            continue
-        else:
-            with Image.open(img_mask_path) as img:
-                width, height = img.size
-
-            if width % 32 != 0 or height % 32 != 0:
-                print(f"Image dimensions must be divisible by 32, got {width}x{height}")
-                continue
-            else:
-                with Image.open(options.img_cond_path) as img_cond:
-                    img_cond_width, img_cond_height = img_cond.size
-
-                if width != img_cond_width or height != img_cond_height:
-                    print(
-                        f"Mask dimensions must match conditioning image, got {width}x{height} and {img_cond_width}x{img_cond_height}"
-                    )
-                    continue
-
-        options.img_mask_path = img_mask_path
-        break
-
-    return options
-
-
-@torch.inference_mode()
-def main(
-    seed: int | None = None,
-    prompt: str = "a white paper cup",
-    device: str = "cuda" if torch.cuda.is_available() else "cpu",
-    num_steps: int = 50,
-    loop: bool = False,
-    guidance: float = 30.0,
-    offload: bool = False,
-    output_dir: str = "output",
-    add_sampling_metadata: bool = True,
-    img_cond_path: str = "assets/cup.png",
-    img_mask_path: str = "assets/cup_mask.png",
-):
-    """
-    Sample the flux model. Either interactively (set `--loop`) or run for a
-    single image. This demo assumes that the conditioning image and mask have
-    the same shape and that height and width are divisible by 32.
-
-    Args:
-        seed: Set a seed for sampling
-        output_name: where to save the output image, `{idx}` will be replaced
-            by the index of the sample
-        prompt: Prompt used for sampling
-        device: Pytorch device
-        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
-        loop: start an interactive session and sample multiple times
-        guidance: guidance value used for guidance distillation
-        add_sampling_metadata: Add the prompt to the image Exif metadata
-        img_cond_path: path to conditioning image (jpeg/png/webp)
-        img_mask_path: path to conditioning mask (jpeg/png/webp
-    """
-    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
-
-    name = "flux-dev-fill"
-    if name not in configs:
-        available = ", ".join(configs.keys())
-        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
-
-    torch_device = torch.device(device)
-
-    output_name = os.path.join(output_dir, "img_{idx}.jpg")
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        idx = 0
-    else:
-        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
-        if len(fns) > 0:
-            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
-        else:
-            idx = 0
-
-    # init all components
-    t5 = load_t5(torch_device, max_length=128)
-    clip = load_clip(torch_device)
-    model = load_flow_model(name, device="cpu" if offload else torch_device)
-    ae = load_ae(name, device="cpu" if offload else torch_device)
-
-    rng = torch.Generator(device="cpu")
-    with Image.open(img_cond_path) as img:
-        width, height = img.size
-    opts = SamplingOptions(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_steps=num_steps,
-        guidance=guidance,
-        seed=seed,
-        img_cond_path=img_cond_path,
-        img_mask_path=img_mask_path,
-    )
-
-    if loop:
-        opts = parse_prompt(opts)
-        opts = parse_img_cond_path(opts)
-
-        with Image.open(opts.img_cond_path) as img:
-            width, height = img.size
-        opts.height = height
-        opts.width = width
-
-        opts = parse_img_mask_path(opts)
-
-    while opts is not None:
-        if opts.seed is None:
-            opts.seed = rng.seed()
-        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
-        t0 = time.perf_counter()
-
-        # prepare input
-        x = get_noise(
-            1,
-            opts.height,
-            opts.width,
-            device=torch_device,
-            dtype=torch.bfloat16,
-            seed=opts.seed,
-        )
-        opts.seed = None
-        if offload:
-            t5, clip, ae = t5.to(torch_device), clip.to(torch_device), ae.to(torch.device)
-        inp = prepare_fill(
-            t5,
-            clip,
-            x,
-            prompt=opts.prompt,
-            ae=ae,
-            img_cond_path=opts.img_cond_path,
-            mask_path=opts.img_mask_path,
-        )
-
-        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
-
-        # offload TEs and AE to CPU, load model to gpu
-        if offload:
-            t5, clip, ae = t5.cpu(), clip.cpu(), ae.cpu()
-            torch.cuda.empty_cache()
-            model = model.to(torch_device)
-
-        # denoise initial noise
-        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
-
-        # offload model, load autoencoder to gpu
-        if offload:
-            model.cpu()
-            torch.cuda.empty_cache()
-            ae.decoder.to(x.device)
-
-        # decode latents to pixel space
-        x = unpack(x.float(), opts.height, opts.width)
-        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
-            x = ae.decode(x)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        t1 = time.perf_counter()
-        print(f"Done in {t1 - t0:.1f}s")
-
-        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
-
-        if loop:
-            print("-" * 80)
-            opts = parse_prompt(opts)
-            opts = parse_img_cond_path(opts)
-
-            with Image.open(opts.img_cond_path) as img:
-                width, height = img.size
-            opts.height = height
-            opts.width = width
-
-            opts = parse_img_mask_path(opts)
-        else:
-            opts = None
-
-
-def app():
-    Fire(main)
-
-
-if __name__ == "__main__":
-    app()

flux/cli_redux.py

@@ -1,279 +0,0 @@
-import os
-import re
-import time
-from dataclasses import dataclass
-from glob import iglob
-
-import torch
-from fire import Fire
-from transformers import pipeline
-
-from flux.modules.image_embedders import ReduxImageEncoder
-from flux.sampling import denoise, get_noise, get_schedule, prepare_redux, unpack
-from flux.util import configs, load_ae, load_clip, load_flow_model, load_t5, save_image
-
-
-@dataclass
-class SamplingOptions:
-    prompt: str
-    width: int
-    height: int
-    num_steps: int
-    guidance: float
-    seed: int | None
-    img_cond_path: str
-
-
-def parse_prompt(options: SamplingOptions) -> SamplingOptions | None:
-    user_question = "Write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Leave this field empty to do nothing "
-        "or write a command starting with a slash:\n"
-        "- '/w <width>' will set the width of the generated image\n"
-        "- '/h <height>' will set the height of the generated image\n"
-        "- '/s <seed>' sets the next seed\n"
-        "- '/g <guidance>' sets the guidance (flux-dev only)\n"
-        "- '/n <steps>' sets the number of steps\n"
-        "- '/q' to quit"
-    )
-
-    while (prompt := input(user_question)).startswith("/"):
-        if prompt.startswith("/w"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, width = prompt.split()
-            options.width = 16 * (int(width) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/h"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, height = prompt.split()
-            options.height = 16 * (int(height) // 16)
-            print(
-                f"Setting resolution to {options.width} x {options.height} "
-                f"({options.height *options.width/1e6:.2f}MP)"
-            )
-        elif prompt.startswith("/g"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, guidance = prompt.split()
-            options.guidance = float(guidance)
-            print(f"Setting guidance to {options.guidance}")
-        elif prompt.startswith("/s"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, seed = prompt.split()
-            options.seed = int(seed)
-            print(f"Setting seed to {options.seed}")
-        elif prompt.startswith("/n"):
-            if prompt.count(" ") != 1:
-                print(f"Got invalid command '{prompt}'\n{usage}")
-                continue
-            _, steps = prompt.split()
-            options.num_steps = int(steps)
-            print(f"Setting number of steps to {options.num_steps}")
-        elif prompt.startswith("/q"):
-            print("Quitting")
-            return None
-        else:
-            if not prompt.startswith("/h"):
-                print(f"Got invalid command '{prompt}'\n{usage}")
-            print(usage)
-    return options
-
-
-def parse_img_cond_path(options: SamplingOptions | None) -> SamplingOptions | None:
-    if options is None:
-        return None
-
-    user_question = "Next conditioning image (write /h for help, /q to quit and leave empty to repeat):\n"
-    usage = (
-        "Usage: Either write your prompt directly, leave this field empty "
-        "to repeat the conditioning image or write a command starting with a slash:\n"
-        "- '/q' to quit"
-    )
-
-    while True:
-        img_cond_path = input(user_question)
-
-        if img_cond_path.startswith("/"):
-            if img_cond_path.startswith("/q"):
-                print("Quitting")
-                return None
-            else:
-                if not img_cond_path.startswith("/h"):
-                    print(f"Got invalid command '{img_cond_path}'\n{usage}")
-                print(usage)
-            continue
-
-        if img_cond_path == "":
-            break
-
-        if not os.path.isfile(img_cond_path) or not img_cond_path.lower().endswith(
-            (".jpg", ".jpeg", ".png", ".webp")
-        ):
-            print(f"File '{img_cond_path}' does not exist or is not a valid image file")
-            continue
-
-        options.img_cond_path = img_cond_path
-        break
-
-    return options
-
-
-@torch.inference_mode()
-def main(
-    name: str = "flux-dev",
-    width: int = 1360,
-    height: int = 768,
-    seed: int | None = None,
-    device: str = "cuda" if torch.cuda.is_available() else "cpu",
-    num_steps: int | None = None,
-    loop: bool = False,
-    guidance: float = 2.5,
-    offload: bool = False,
-    output_dir: str = "output",
-    add_sampling_metadata: bool = True,
-    img_cond_path: str = "assets/robot.webp",
-):
-    """
-    Sample the flux model. Either interactively (set `--loop`) or run for a
-    single image.
-
-    Args:
-        name: Name of the model to load
-        height: height of the sample in pixels (should be a multiple of 16)
-        width: width of the sample in pixels (should be a multiple of 16)
-        seed: Set a seed for sampling
-        output_name: where to save the output image, `{idx}` will be replaced
-            by the index of the sample
-        prompt: Prompt used for sampling
-        device: Pytorch device
-        num_steps: number of sampling steps (default 4 for schnell, 50 for guidance distilled)
-        loop: start an interactive session and sample multiple times
-        guidance: guidance value used for guidance distillation
-        add_sampling_metadata: Add the prompt to the image Exif metadata
-        img_cond_path: path to conditioning image (jpeg/png/webp)
-    """
-    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
-
-    if name not in configs:
-        available = ", ".join(configs.keys())
-        raise ValueError(f"Got unknown model name: {name}, chose from {available}")
-
-    torch_device = torch.device(device)
-    if num_steps is None:
-        num_steps = 4 if name == "flux-schnell" else 50
-
-    output_name = os.path.join(output_dir, "img_{idx}.jpg")
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        idx = 0
-    else:
-        fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)]
-        if len(fns) > 0:
-            idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1
-        else:
-            idx = 0
-
-    # init all components
-    t5 = load_t5(torch_device, max_length=256 if name == "flux-schnell" else 512)
-    clip = load_clip(torch_device)
-    model = load_flow_model(name, device="cpu" if offload else torch_device)
-    ae = load_ae(name, device="cpu" if offload else torch_device)
-    img_embedder = ReduxImageEncoder(torch_device)
-
-    rng = torch.Generator(device="cpu")
-    prompt = ""
-    opts = SamplingOptions(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_steps=num_steps,
-        guidance=guidance,
-        seed=seed,
-        img_cond_path=img_cond_path,
-    )
-
-    if loop:
-        opts = parse_prompt(opts)
-        opts = parse_img_cond_path(opts)
-
-    while opts is not None:
-        if opts.seed is None:
-            opts.seed = rng.seed()
-        print(f"Generating with seed {opts.seed}:\n{opts.prompt}")
-        t0 = time.perf_counter()
-
-        # prepare input
-        x = get_noise(
-            1,
-            opts.height,
-            opts.width,
-            device=torch_device,
-            dtype=torch.bfloat16,
-            seed=opts.seed,
-        )
-        opts.seed = None
-        if offload:
-            ae = ae.cpu()
-            torch.cuda.empty_cache()
-            t5, clip = t5.to(torch_device), clip.to(torch_device)
-        inp = prepare_redux(
-            t5,
-            clip,
-            x,
-            prompt=opts.prompt,
-            encoder=img_embedder,
-            img_cond_path=opts.img_cond_path,
-        )
-        timesteps = get_schedule(opts.num_steps, inp["img"].shape[1], shift=(name != "flux-schnell"))
-
-        # offload TEs to CPU, load model to gpu
-        if offload:
-            t5, clip = t5.cpu(), clip.cpu()
-            torch.cuda.empty_cache()
-            model = model.to(torch_device)
-
-        # denoise initial noise
-        x = denoise(model, **inp, timesteps=timesteps, guidance=opts.guidance)
-
-        # offload model, load autoencoder to gpu
-        if offload:
-            model.cpu()
-            torch.cuda.empty_cache()
-            ae.decoder.to(x.device)
-
-        # decode latents to pixel space
-        x = unpack(x.float(), opts.height, opts.width)
-        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
-            x = ae.decode(x)
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        t1 = time.perf_counter()
-        print(f"Done in {t1 - t0:.1f}s")
-
-        idx = save_image(nsfw_classifier, name, output_name, idx, x, add_sampling_metadata, prompt)
-
-        if loop:
-            print("-" * 80)
-            opts = parse_prompt(opts)
-            opts = parse_img_cond_path(opts)
-        else:
-            opts = None
-
-
-def app():
-    Fire(main)
-
-
-if __name__ == "__main__":
-    app()

flux/math.py

@@ -1,30 +0,0 @@
-import torch
-from einops import rearrange
-from torch import Tensor
-
-
-def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
-    q, k = apply_rope(q, k, pe)
-
-    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
-    x = rearrange(x, "B H L D -> B L (H D)")
-
-    return x
-
-
-def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
-    assert dim % 2 == 0
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
-    omega = 1.0 / (theta**scale)
-    out = torch.einsum("...n,d->...nd", pos, omega)
-    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
-    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
-    return out.float()
-
-
-def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
-    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
-    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
-    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
-    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
-    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

flux/model.py

@@ -1,143 +0,0 @@
-from dataclasses import dataclass
-
-import torch
-from torch import Tensor, nn
-
-from flux.modules.layers import (
-    DoubleStreamBlock,
-    EmbedND,
-    LastLayer,
-    MLPEmbedder,
-    SingleStreamBlock,
-    timestep_embedding,
-)
-from flux.modules.lora import LinearLora, replace_linear_with_lora
-
-
-@dataclass
-class FluxParams:
-    in_channels: int
-    out_channels: int
-    vec_in_dim: int
-    context_in_dim: int
-    hidden_size: int
-    mlp_ratio: float
-    num_heads: int
-    depth: int
-    depth_single_blocks: int
-    axes_dim: list[int]
-    theta: int
-    qkv_bias: bool
-    guidance_embed: bool
-
-
-class Flux(nn.Module):
-    """
-    Transformer model for flow matching on sequences.
-    """
-
-    def __init__(self, params: FluxParams):
-        super().__init__()
-
-        self.params = params
-        self.in_channels = params.in_channels
-        self.out_channels = params.out_channels
-        if params.hidden_size % params.num_heads != 0:
-            raise ValueError(
-                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
-            )
-        pe_dim = params.hidden_size // params.num_heads
-        if sum(params.axes_dim) != pe_dim:
-            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
-        self.hidden_size = params.hidden_size
-        self.num_heads = params.num_heads
-        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
-        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
-        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
-        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
-        self.guidance_in = (
-            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
-        )
-        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
-
-        self.double_blocks = nn.ModuleList(
-            [
-                DoubleStreamBlock(
-                    self.hidden_size,
-                    self.num_heads,
-                    mlp_ratio=params.mlp_ratio,
-                    qkv_bias=params.qkv_bias,
-                )
-                for _ in range(params.depth)
-            ]
-        )
-
-        self.single_blocks = nn.ModuleList(
-            [
-                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
-                for _ in range(params.depth_single_blocks)
-            ]
-        )
-
-        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
-
-    def forward(
-        self,
-        img: Tensor,
-        img_ids: Tensor,
-        txt: Tensor,
-        txt_ids: Tensor,
-        timesteps: Tensor,
-        y: Tensor,
-        guidance: Tensor | None = None,
-    ) -> Tensor:
-        if img.ndim != 3 or txt.ndim != 3:
-            raise ValueError("Input img and txt tensors must have 3 dimensions.")
-
-        # running on sequences img
-        img = self.img_in(img)
-        vec = self.time_in(timestep_embedding(timesteps, 256))
-        if self.params.guidance_embed:
-            if guidance is None:
-                raise ValueError("Didn't get guidance strength for guidance distilled model.")
-            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
-        vec = vec + self.vector_in(y)
-        txt = self.txt_in(txt)
-
-        ids = torch.cat((txt_ids, img_ids), dim=1)
-        pe = self.pe_embedder(ids)
-
-        for block in self.double_blocks:
-            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
-
-        img = torch.cat((txt, img), 1)
-        for block in self.single_blocks:
-            img = block(img, vec=vec, pe=pe)
-        img = img[:, txt.shape[1] :, ...]
-
-        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
-        return img
-
-
-class FluxLoraWrapper(Flux):
-    def __init__(
-        self,
-        lora_rank: int = 128,
-        lora_scale: float = 1.0,
-        *args,
-        **kwargs,
-    ) -> None:
-        super().__init__(*args, **kwargs)
-
-        self.lora_rank = lora_rank
-
-        replace_linear_with_lora(
-            self,
-            max_rank=lora_rank,
-            scale=lora_scale,
-        )
-
-    def set_lora_scale(self, scale: float) -> None:
-        for module in self.modules():
-            if isinstance(module, LinearLora):
-                module.set_scale(scale=scale)

flux/modules/autoencoder.py

@@ -1,312 +0,0 @@
-from dataclasses import dataclass
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-
-
-@dataclass
-class AutoEncoderParams:
-    resolution: int
-    in_channels: int
-    ch: int
-    out_ch: int
-    ch_mult: list[int]
-    num_res_blocks: int
-    z_channels: int
-    scale_factor: float
-    shift_factor: float
-
-
-def swish(x: Tensor) -> Tensor:
-    return x * torch.sigmoid(x)
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-
-    def attention(self, h_: Tensor) -> Tensor:
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        b, c, h, w = q.shape
-        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
-        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
-        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
-        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
-
-        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x + self.proj_out(self.attention(x))
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-
-        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        if self.in_channels != self.out_channels:
-            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = swish(h)
-        h = self.conv1(h)
-
-        h = self.norm2(h)
-        h = swish(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            x = self.nin_shortcut(x)
-
-        return x + h
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        # no asymmetric padding in torch conv, must do it ourselves
-        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-
-    def forward(self, x: Tensor):
-        pad = (0, 1, 0, 1)
-        x = nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x: Tensor):
-        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        x = self.conv(x)
-        return x
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: int,
-        ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        # downsampling
-        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,) + tuple(ch_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        block_in = self.ch
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch * in_ch_mult[i_level]
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                down.downsample = Downsample(block_in)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
-        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x: Tensor) -> Tensor:
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1])
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions - 1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-        # end
-        h = self.norm_out(h)
-        h = swish(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        ch: int,
-        out_ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        in_channels: int,
-        resolution: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.ffactor = 2 ** (self.num_resolutions - 1)
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.z_shape = (1, z_channels, curr_res, curr_res)
-
-        # z to block_in
-        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks + 1):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
-        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, z: Tensor) -> Tensor:
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.up[i_level].block[i_block](h)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = swish(h)
-        h = self.conv_out(h)
-        return h
-
-
-class DiagonalGaussian(nn.Module):
-    def __init__(self, sample: bool = True, chunk_dim: int = 1):
-        super().__init__()
-        self.sample = sample
-        self.chunk_dim = chunk_dim
-
-    def forward(self, z: Tensor) -> Tensor:
-        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
-        if self.sample:
-            std = torch.exp(0.5 * logvar)
-            return mean + std * torch.randn_like(mean)
-        else:
-            return mean
-
-
-class AutoEncoder(nn.Module):
-    def __init__(self, params: AutoEncoderParams):
-        super().__init__()
-        self.encoder = Encoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=params.ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-        self.decoder = Decoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=params.ch,
-            out_ch=params.out_ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-        self.reg = DiagonalGaussian()
-
-        self.scale_factor = params.scale_factor
-        self.shift_factor = params.shift_factor
-
-    def encode(self, x: Tensor) -> Tensor:
-        z = self.reg(self.encoder(x))
-        z = self.scale_factor * (z - self.shift_factor)
-        return z
-
-    def decode(self, z: Tensor) -> Tensor:
-        z = z / self.scale_factor + self.shift_factor
-        return self.decoder(z)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.decode(self.encode(x))

flux/modules/conditioner.py

@@ -1,37 +0,0 @@
-from torch import Tensor, nn
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
-
-
-class HFEmbedder(nn.Module):
-    def __init__(self, version: str, max_length: int, **hf_kwargs):
-        super().__init__()
-        self.is_clip = version.startswith("openai")
-        self.max_length = max_length
-        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
-
-        if self.is_clip:
-            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
-            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
-        else:
-            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
-            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
-
-        self.hf_module = self.hf_module.eval().requires_grad_(False)
-
-    def forward(self, text: list[str]) -> Tensor:
-        batch_encoding = self.tokenizer(
-            text,
-            truncation=True,
-            max_length=self.max_length,
-            return_length=False,
-            return_overflowing_tokens=False,
-            padding="max_length",
-            return_tensors="pt",
-        )
-
-        outputs = self.hf_module(
-            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
-            attention_mask=None,
-            output_hidden_states=False,
-        )
-        return outputs[self.output_key]

flux/modules/image_embedders.py

@@ -1,103 +0,0 @@
-import os
-
-import cv2
-import numpy as np
-import torch
-from einops import rearrange, repeat
-from PIL import Image
-from safetensors.torch import load_file as load_sft
-from torch import nn
-from transformers import AutoModelForDepthEstimation, AutoProcessor, SiglipImageProcessor, SiglipVisionModel
-
-from flux.util import print_load_warning
-
-
-class DepthImageEncoder:
-    depth_model_name = "LiheYoung/depth-anything-large-hf"
-
-    def __init__(self, device):
-        self.device = device
-        self.depth_model = AutoModelForDepthEstimation.from_pretrained(self.depth_model_name).to(device)
-        self.processor = AutoProcessor.from_pretrained(self.depth_model_name)
-
-    def __call__(self, img: torch.Tensor) -> torch.Tensor:
-        hw = img.shape[-2:]
-
-        img = torch.clamp(img, -1.0, 1.0)
-        img_byte = ((img + 1.0) * 127.5).byte()
-
-        img = self.processor(img_byte, return_tensors="pt")["pixel_values"]
-        depth = self.depth_model(img.to(self.device)).predicted_depth
-        depth = repeat(depth, "b h w -> b 3 h w")
-        depth = torch.nn.functional.interpolate(depth, hw, mode="bicubic", antialias=True)
-
-        depth = depth / 127.5 - 1.0
-        return depth
-
-
-class CannyImageEncoder:
-    def __init__(
-        self,
-        device,
-        min_t: int = 50,
-        max_t: int = 200,
-    ):
-        self.device = device
-        self.min_t = min_t
-        self.max_t = max_t
-
-    def __call__(self, img: torch.Tensor) -> torch.Tensor:
-        assert img.shape[0] == 1, "Only batch size 1 is supported"
-
-        img = rearrange(img[0], "c h w -> h w c")
-        img = torch.clamp(img, -1.0, 1.0)
-        img_np = ((img + 1.0) * 127.5).numpy().astype(np.uint8)
-
-        # Apply Canny edge detection
-        canny = cv2.Canny(img_np, self.min_t, self.max_t)
-
-        # Convert back to torch tensor and reshape
-        canny = torch.from_numpy(canny).float() / 127.5 - 1.0
-        canny = rearrange(canny, "h w -> 1 1 h w")
-        canny = repeat(canny, "b 1 ... -> b 3 ...")
-        return canny.to(self.device)
-
-
-class ReduxImageEncoder(nn.Module):
-    siglip_model_name = "google/siglip-so400m-patch14-384"
-
-    def __init__(
-        self,
-        device,
-        redux_dim: int = 1152,
-        txt_in_features: int = 4096,
-        redux_path: str | None = os.getenv("FLUX_REDUX"),
-        dtype=torch.bfloat16,
-    ) -> None:
-        assert redux_path is not None, "Redux path must be provided"
-
-        super().__init__()
-
-        self.redux_dim = redux_dim
-        self.device = device if isinstance(device, torch.device) else torch.device(device)
-        self.dtype = dtype
-
-        with self.device:
-            self.redux_up = nn.Linear(redux_dim, txt_in_features * 3, dtype=dtype)
-            self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features, dtype=dtype)
-
-            sd = load_sft(redux_path, device=str(device))
-            missing, unexpected = self.load_state_dict(sd, strict=False, assign=True)
-            print_load_warning(missing, unexpected)
-
-            self.siglip = SiglipVisionModel.from_pretrained(self.siglip_model_name).to(dtype=dtype)
-        self.normalize = SiglipImageProcessor.from_pretrained(self.siglip_model_name)
-
-    def __call__(self, x: Image.Image) -> torch.Tensor:
-        imgs = self.normalize.preprocess(images=[x], do_resize=True, return_tensors="pt", do_convert_rgb=True)
-
-        _encoded_x = self.siglip(**imgs.to(device=self.device, dtype=self.dtype)).last_hidden_state
-
-        projected_x = self.redux_down(nn.functional.silu(self.redux_up(_encoded_x)))
-
-        return projected_x

flux/modules/layers.py

@@ -1,253 +0,0 @@
-import math
-from dataclasses import dataclass
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-
-from flux.math import attention, rope
-
-
-class EmbedND(nn.Module):
-    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
-        super().__init__()
-        self.dim = dim
-        self.theta = theta
-        self.axes_dim = axes_dim
-
-    def forward(self, ids: Tensor) -> Tensor:
-        n_axes = ids.shape[-1]
-        emb = torch.cat(
-            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
-            dim=-3,
-        )
-
-        return emb.unsqueeze(1)
-
-
-def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
-    """
-    Create sinusoidal timestep embeddings.
-    :param t: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an (N, D) Tensor of positional embeddings.
-    """
-    t = time_factor * t
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-        t.device
-    )
-
-    args = t[:, None].float() * freqs[None]
-    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    if torch.is_floating_point(t):
-        embedding = embedding.to(t)
-    return embedding
-
-
-class MLPEmbedder(nn.Module):
-    def __init__(self, in_dim: int, hidden_dim: int):
-        super().__init__()
-        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
-        self.silu = nn.SiLU()
-        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.out_layer(self.silu(self.in_layer(x)))
-
-
-class RMSNorm(torch.nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.scale = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: Tensor):
-        x_dtype = x.dtype
-        x = x.float()
-        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
-        return (x * rrms).to(dtype=x_dtype) * self.scale
-
-
-class QKNorm(torch.nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.query_norm = RMSNorm(dim)
-        self.key_norm = RMSNorm(dim)
-
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
-        q = self.query_norm(q)
-        k = self.key_norm(k)
-        return q.to(v), k.to(v)
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.norm = QKNorm(head_dim)
-        self.proj = nn.Linear(dim, dim)
-
-    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
-        qkv = self.qkv(x)
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        q, k = self.norm(q, k, v)
-        x = attention(q, k, v, pe=pe)
-        x = self.proj(x)
-        return x
-
-
-@dataclass
-class ModulationOut:
-    shift: Tensor
-    scale: Tensor
-    gate: Tensor
-
-
-class Modulation(nn.Module):
-    def __init__(self, dim: int, double: bool):
-        super().__init__()
-        self.is_double = double
-        self.multiplier = 6 if double else 3
-        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
-
-    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
-        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
-
-        return (
-            ModulationOut(*out[:3]),
-            ModulationOut(*out[3:]) if self.is_double else None,
-        )
-
-
-class DoubleStreamBlock(nn.Module):
-    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
-        super().__init__()
-
-        mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        self.num_heads = num_heads
-        self.hidden_size = hidden_size
-        self.img_mod = Modulation(hidden_size, double=True)
-        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.img_mlp = nn.Sequential(
-            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
-            nn.GELU(approximate="tanh"),
-            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
-        )
-
-        self.txt_mod = Modulation(hidden_size, double=True)
-        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.txt_mlp = nn.Sequential(
-            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
-            nn.GELU(approximate="tanh"),
-            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
-        )
-
-    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
-        img_mod1, img_mod2 = self.img_mod(vec)
-        txt_mod1, txt_mod2 = self.txt_mod(vec)
-
-        # prepare image for attention
-        img_modulated = self.img_norm1(img)
-        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
-        img_qkv = self.img_attn.qkv(img_modulated)
-        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
-
-        # prepare txt for attention
-        txt_modulated = self.txt_norm1(txt)
-        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
-        txt_qkv = self.txt_attn.qkv(txt_modulated)
-        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
-
-        # run actual attention
-        q = torch.cat((txt_q, img_q), dim=2)
-        k = torch.cat((txt_k, img_k), dim=2)
-        v = torch.cat((txt_v, img_v), dim=2)
-
-        attn = attention(q, k, v, pe=pe)
-        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
-
-        # calculate the img bloks
-        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
-        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
-
-        # calculate the txt bloks
-        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
-        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
-        return img, txt
-
-
-class SingleStreamBlock(nn.Module):
-    """
-    A DiT block with parallel linear layers as described in
-    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        qk_scale: float | None = None,
-    ):
-        super().__init__()
-        self.hidden_dim = hidden_size
-        self.num_heads = num_heads
-        head_dim = hidden_size // num_heads
-        self.scale = qk_scale or head_dim**-0.5
-
-        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        # qkv and mlp_in
-        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
-        # proj and mlp_out
-        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
-
-        self.norm = QKNorm(head_dim)
-
-        self.hidden_size = hidden_size
-        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-
-        self.mlp_act = nn.GELU(approximate="tanh")
-        self.modulation = Modulation(hidden_size, double=False)
-
-    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
-        mod, _ = self.modulation(vec)
-        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
-        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
-
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        q, k = self.norm(q, k, v)
-
-        # compute attention
-        attn = attention(q, k, v, pe=pe)
-        # compute activation in mlp stream, cat again and run second linear layer
-        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
-        return x + mod.gate * output
-
-
-class LastLayer(nn.Module):
-    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
-
-    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
-        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
-        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
-        x = self.linear(x)
-        return x

flux/modules/lora.py

@@ -1,94 +0,0 @@
-import torch
-from torch import nn
-
-
-def replace_linear_with_lora(
-    module: nn.Module,
-    max_rank: int,
-    scale: float = 1.0,
-) -> None:
-    for name, child in module.named_children():
-        if isinstance(child, nn.Linear):
-            new_lora = LinearLora(
-                in_features=child.in_features,
-                out_features=child.out_features,
-                bias=child.bias,
-                rank=max_rank,
-                scale=scale,
-                dtype=child.weight.dtype,
-                device=child.weight.device,
-            )
-
-            new_lora.weight = child.weight
-            new_lora.bias = child.bias if child.bias is not None else None
-
-            setattr(module, name, new_lora)
-        else:
-            replace_linear_with_lora(
-                module=child,
-                max_rank=max_rank,
-                scale=scale,
-            )
-
-
-class LinearLora(nn.Linear):
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias: bool,
-        rank: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        lora_bias: bool = True,
-        scale: float = 1.0,
-        *args,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            in_features=in_features,
-            out_features=out_features,
-            bias=bias is not None,
-            device=device,
-            dtype=dtype,
-            *args,
-            **kwargs,
-        )
-
-        assert isinstance(scale, float), "scale must be a float"
-
-        self.scale = scale
-        self.rank = rank
-        self.lora_bias = lora_bias
-        self.dtype = dtype
-        self.device = device
-
-        if rank > (new_rank := min(self.out_features, self.in_features)):
-            self.rank = new_rank
-
-        self.lora_A = nn.Linear(
-            in_features=in_features,
-            out_features=self.rank,
-            bias=False,
-            dtype=dtype,
-            device=device,
-        )
-        self.lora_B = nn.Linear(
-            in_features=self.rank,
-            out_features=out_features,
-            bias=self.lora_bias,
-            dtype=dtype,
-            device=device,
-        )
-
-    def set_scale(self, scale: float) -> None:
-        assert isinstance(scale, float), "scalar value must be a float"
-        self.scale = scale
-
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        base_out = super().forward(input)
-
-        _lora_out_B = self.lora_B(self.lora_A(input))
-        lora_update = _lora_out_B * self.scale
-
-        return base_out + lora_update

flux/sampling.py

@@ -1,282 +0,0 @@
-import math
-from typing import Callable
-
-import numpy as np
-import torch
-from einops import rearrange, repeat
-from PIL import Image
-from torch import Tensor
-
-from .model import Flux
-from .modules.autoencoder import AutoEncoder
-from .modules.conditioner import HFEmbedder
-from .modules.image_embedders import CannyImageEncoder, DepthImageEncoder, ReduxImageEncoder
-
-
-def get_noise(
-    num_samples: int,
-    height: int,
-    width: int,
-    device: torch.device,
-    dtype: torch.dtype,
-    seed: int,
-):
-    return torch.randn(
-        num_samples,
-        16,
-        # allow for packing
-        2 * math.ceil(height / 16),
-        2 * math.ceil(width / 16),
-        device=device,
-        dtype=dtype,
-        generator=torch.Generator(device=device).manual_seed(seed),
-    )
-
-
-def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
-    bs, c, h, w = img.shape
-    if bs == 1 and not isinstance(prompt, str):
-        bs = len(prompt)
-
-    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-    if img.shape[0] == 1 and bs > 1:
-        img = repeat(img, "1 ... -> bs ...", bs=bs)
-
-    img_ids = torch.zeros(h // 2, w // 2, 3)
-    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
-    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
-    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
-
-    if isinstance(prompt, str):
-        prompt = [prompt]
-    txt = t5(prompt)
-    if txt.shape[0] == 1 and bs > 1:
-        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
-    txt_ids = torch.zeros(bs, txt.shape[1], 3)
-
-    vec = clip(prompt)
-    if vec.shape[0] == 1 and bs > 1:
-        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
-
-    return {
-        "img": img,
-        "img_ids": img_ids.to(img.device),
-        "txt": txt.to(img.device),
-        "txt_ids": txt_ids.to(img.device),
-        "vec": vec.to(img.device),
-    }
-
-
-def prepare_control(
-    t5: HFEmbedder,
-    clip: HFEmbedder,
-    img: Tensor,
-    prompt: str | list[str],
-    ae: AutoEncoder,
-    encoder: DepthImageEncoder | CannyImageEncoder,
-    img_cond_path: str,
-) -> dict[str, Tensor]:
-    # load and encode the conditioning image
-    bs, _, h, w = img.shape
-    if bs == 1 and not isinstance(prompt, str):
-        bs = len(prompt)
-
-    img_cond = Image.open(img_cond_path).convert("RGB")
-
-    width = w * 8
-    height = h * 8
-    img_cond = img_cond.resize((width, height), Image.LANCZOS)
-    img_cond = np.array(img_cond)
-    img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
-    img_cond = rearrange(img_cond, "h w c -> 1 c h w")
-
-    with torch.no_grad():
-        img_cond = encoder(img_cond)
-        img_cond = ae.encode(img_cond)
-
-    img_cond = img_cond.to(torch.bfloat16)
-    img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-    if img_cond.shape[0] == 1 and bs > 1:
-        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
-
-    return_dict = prepare(t5, clip, img, prompt)
-    return_dict["img_cond"] = img_cond
-    return return_dict
-
-
-def prepare_fill(
-    t5: HFEmbedder,
-    clip: HFEmbedder,
-    img: Tensor,
-    prompt: str | list[str],
-    ae: AutoEncoder,
-    img_cond_path: str,
-    mask_path: str,
-) -> dict[str, Tensor]:
-    # load and encode the conditioning image and the mask
-    bs, _, _, _ = img.shape
-    if bs == 1 and not isinstance(prompt, str):
-        bs = len(prompt)
-
-    img_cond = Image.open(img_cond_path).convert("RGB")
-    img_cond = np.array(img_cond)
-    img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
-    img_cond = rearrange(img_cond, "h w c -> 1 c h w")
-
-    mask = Image.open(mask_path).convert("L")
-    mask = np.array(mask)
-    mask = torch.from_numpy(mask).float() / 255.0
-    mask = rearrange(mask, "h w -> 1 1 h w")
-
-    with torch.no_grad():
-        img_cond = img_cond.to(img.device)
-        mask = mask.to(img.device)
-        img_cond = img_cond * (1 - mask)
-        img_cond = ae.encode(img_cond)
-        mask = mask[:, 0, :, :]
-        mask = mask.to(torch.bfloat16)
-        mask = rearrange(
-            mask,
-            "b (h ph) (w pw) -> b (ph pw) h w",
-            ph=8,
-            pw=8,
-        )
-        mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-        if mask.shape[0] == 1 and bs > 1:
-            mask = repeat(mask, "1 ... -> bs ...", bs=bs)
-
-    img_cond = img_cond.to(torch.bfloat16)
-    img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-    if img_cond.shape[0] == 1 and bs > 1:
-        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
-
-    img_cond = torch.cat((img_cond, mask), dim=-1)
-
-    return_dict = prepare(t5, clip, img, prompt)
-    return_dict["img_cond"] = img_cond.to(img.device)
-    return return_dict
-
-
-def prepare_redux(
-    t5: HFEmbedder,
-    clip: HFEmbedder,
-    img: Tensor,
-    prompt: str | list[str],
-    encoder: ReduxImageEncoder,
-    img_cond_path: str,
-) -> dict[str, Tensor]:
-    bs, _, h, w = img.shape
-    if bs == 1 and not isinstance(prompt, str):
-        bs = len(prompt)
-
-    img_cond = Image.open(img_cond_path).convert("RGB")
-    with torch.no_grad():
-        img_cond = encoder(img_cond)
-
-    img_cond = img_cond.to(torch.bfloat16)
-    if img_cond.shape[0] == 1 and bs > 1:
-        img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
-
-    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-    if img.shape[0] == 1 and bs > 1:
-        img = repeat(img, "1 ... -> bs ...", bs=bs)
-
-    img_ids = torch.zeros(h // 2, w // 2, 3)
-    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
-    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
-    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
-
-    if isinstance(prompt, str):
-        prompt = [prompt]
-    txt = t5(prompt)
-    txt = torch.cat((txt, img_cond.to(txt)), dim=-2)
-    if txt.shape[0] == 1 and bs > 1:
-        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
-    txt_ids = torch.zeros(bs, txt.shape[1], 3)
-
-    vec = clip(prompt)
-    if vec.shape[0] == 1 and bs > 1:
-        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
-
-    return {
-        "img": img,
-        "img_ids": img_ids.to(img.device),
-        "txt": txt.to(img.device),
-        "txt_ids": txt_ids.to(img.device),
-        "vec": vec.to(img.device),
-    }
-
-
-def time_shift(mu: float, sigma: float, t: Tensor):
-    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
-
-
-def get_lin_function(
-    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
-) -> Callable[[float], float]:
-    m = (y2 - y1) / (x2 - x1)
-    b = y1 - m * x1
-    return lambda x: m * x + b
-
-
-def get_schedule(
-    num_steps: int,
-    image_seq_len: int,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-    shift: bool = True,
-) -> list[float]:
-    # extra step for zero
-    timesteps = torch.linspace(1, 0, num_steps + 1)
-
-    # shifting the schedule to favor high timesteps for higher signal images
-    if shift:
-        # estimate mu based on linear estimation between two points
-        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
-        timesteps = time_shift(mu, 1.0, timesteps)
-
-    return timesteps.tolist()
-
-
-def denoise(
-    model: Flux,
-    # model input
-    img: Tensor,
-    img_ids: Tensor,
-    txt: Tensor,
-    txt_ids: Tensor,
-    vec: Tensor,
-    # sampling parameters
-    timesteps: list[float],
-    guidance: float = 4.0,
-    # extra img tokens
-    img_cond: Tensor | None = None,
-):
-    # this is ignored for schnell
-    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
-    for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):
-        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
-        pred = model(
-            img=torch.cat((img, img_cond), dim=-1) if img_cond is not None else img,
-            img_ids=img_ids,
-            txt=txt,
-            txt_ids=txt_ids,
-            y=vec,
-            timesteps=t_vec,
-            guidance=guidance_vec,
-        )
-
-        img = img + (t_prev - t_curr) * pred
-
-    return img
-
-
-def unpack(x: Tensor, height: int, width: int) -> Tensor:
-    return rearrange(
-        x,
-        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
-        h=math.ceil(height / 16),
-        w=math.ceil(width / 16),
-        ph=2,
-        pw=2,
-    )

flux/util.py

@@ -1,447 +0,0 @@
-import os
-from dataclasses import dataclass
-
-import torch
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-from imwatermark import WatermarkEncoder
-from PIL import ExifTags, Image
-from safetensors.torch import load_file as load_sft
-
-from flux.model import Flux, FluxLoraWrapper, FluxParams
-from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
-from flux.modules.conditioner import HFEmbedder
-
-
-def save_image(
-    nsfw_classifier,
-    name: str,
-    output_name: str,
-    idx: int,
-    x: torch.Tensor,
-    add_sampling_metadata: bool,
-    prompt: str,
-    nsfw_threshold: float = 0.85,
-) -> int:
-    fn = output_name.format(idx=idx)
-    print(f"Saving {fn}")
-    # bring into PIL format and save
-    x = x.clamp(-1, 1)
-    x = embed_watermark(x.float())
-    x = rearrange(x[0], "c h w -> h w c")
-
-    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
-    nsfw_score = [x["score"] for x in nsfw_classifier(img) if x["label"] == "nsfw"][0]
-
-    if nsfw_score < nsfw_threshold:
-        exif_data = Image.Exif()
-        exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
-        exif_data[ExifTags.Base.Make] = "Black Forest Labs"
-        exif_data[ExifTags.Base.Model] = name
-        if add_sampling_metadata:
-            exif_data[ExifTags.Base.ImageDescription] = prompt
-        img.save(fn, exif=exif_data, quality=95, subsampling=0)
-        idx += 1
-    else:
-        print("Your generated image may contain NSFW content.")
-
-    return idx
-
-
-@dataclass
-class ModelSpec:
-    params: FluxParams
-    ae_params: AutoEncoderParams
-    ckpt_path: str | None
-    lora_path: str | None
-    ae_path: str | None
-    repo_id: str | None
-    repo_flow: str | None
-    repo_ae: str | None
-
-
-configs = {
-    "flux-dev": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-dev",
-        repo_flow="flux1-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV"),
-        lora_path=None,
-        params=FluxParams(
-            in_channels=64,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-schnell": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-schnell",
-        repo_flow="flux1-schnell.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_SCHNELL"),
-        lora_path=None,
-        params=FluxParams(
-            in_channels=64,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=False,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-dev-canny": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-Canny-dev",
-        repo_flow="flux1-canny-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV_CANNY"),
-        lora_path=None,
-        params=FluxParams(
-            in_channels=128,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-dev-canny-lora": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-dev",
-        repo_flow="flux1-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV"),
-        lora_path=os.getenv("FLUX_DEV_CANNY_LORA"),
-        params=FluxParams(
-            in_channels=128,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-dev-depth": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-Depth-dev",
-        repo_flow="flux1-depth-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV_DEPTH"),
-        lora_path=None,
-        params=FluxParams(
-            in_channels=128,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-dev-depth-lora": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-dev",
-        repo_flow="flux1-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV"),
-        lora_path=os.getenv("FLUX_DEV_DEPTH_LORA"),
-        params=FluxParams(
-            in_channels=128,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-    "flux-dev-fill": ModelSpec(
-        repo_id="black-forest-labs/FLUX.1-Fill-dev",
-        repo_flow="flux1-fill-dev.safetensors",
-        repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_DEV_FILL"),
-        lora_path=None,
-        params=FluxParams(
-            in_channels=384,
-            out_channels=64,
-            vec_in_dim=768,
-            context_in_dim=4096,
-            hidden_size=3072,
-            mlp_ratio=4.0,
-            num_heads=24,
-            depth=19,
-            depth_single_blocks=38,
-            axes_dim=[16, 56, 56],
-            theta=10_000,
-            qkv_bias=True,
-            guidance_embed=True,
-        ),
-        ae_path=os.getenv("AE"),
-        ae_params=AutoEncoderParams(
-            resolution=256,
-            in_channels=3,
-            ch=128,
-            out_ch=3,
-            ch_mult=[1, 2, 4, 4],
-            num_res_blocks=2,
-            z_channels=16,
-            scale_factor=0.3611,
-            shift_factor=0.1159,
-        ),
-    ),
-}
-
-
-def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
-    if len(missing) > 0 and len(unexpected) > 0:
-        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
-        print("\n" + "-" * 79 + "\n")
-        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
-    elif len(missing) > 0:
-        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
-    elif len(unexpected) > 0:
-        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
-
-
-def load_flow_model(
-    name: str, device: str | torch.device = "cuda", hf_download: bool = True, verbose: bool = False
-) -> Flux:
-    # Loading Flux
-    print("Init model")
-    ckpt_path = configs[name].ckpt_path
-    lora_path = configs[name].lora_path
-    if (
-        ckpt_path is None
-        and configs[name].repo_id is not None
-        and configs[name].repo_flow is not None
-        and hf_download
-    ):
-        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
-
-    with torch.device("meta" if ckpt_path is not None else device):
-        if lora_path is not None:
-            model = FluxLoraWrapper(params=configs[name].params).to(torch.bfloat16)
-        else:
-            model = Flux(configs[name].params).to(torch.bfloat16)
-
-    if ckpt_path is not None:
-        print("Loading checkpoint")
-        # load_sft doesn't support torch.device
-        sd = load_sft(ckpt_path, device=str(device))
-        sd = optionally_expand_state_dict(model, sd)
-        missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
-        if verbose:
-            print_load_warning(missing, unexpected)
-
-    if configs[name].lora_path is not None:
-        print("Loading LoRA")
-        lora_sd = load_sft(configs[name].lora_path, device=str(device))
-        # loading the lora params + overwriting scale values in the norms
-        missing, unexpected = model.load_state_dict(lora_sd, strict=False, assign=True)
-        if verbose:
-            print_load_warning(missing, unexpected)
-    return model
-
-
-def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
-    # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
-    return HFEmbedder("google/t5-v1_1-xxl", max_length=max_length, torch_dtype=torch.bfloat16).to(device)
-
-
-def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
-    return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
-
-
-def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
-    ckpt_path = configs[name].ae_path
-    if (
-        ckpt_path is None
-        and configs[name].repo_id is not None
-        and configs[name].repo_ae is not None
-        and hf_download
-    ):
-        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
-
-    # Loading the autoencoder
-    print("Init AE")
-    with torch.device("meta" if ckpt_path is not None else device):
-        ae = AutoEncoder(configs[name].ae_params)
-
-    if ckpt_path is not None:
-        sd = load_sft(ckpt_path, device=str(device))
-        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
-        print_load_warning(missing, unexpected)
-    return ae
-
-
-def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> dict:
-    """
-    Optionally expand the state dict to match the model's parameters shapes.
-    """
-    for name, param in model.named_parameters():
-        if name in state_dict:
-            if state_dict[name].shape != param.shape:
-                print(
-                    f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape {param.shape}."
-                )
-                # expand with zeros:
-                expanded_state_dict_weight = torch.zeros_like(param, device=state_dict[name].device)
-                slices = tuple(slice(0, dim) for dim in state_dict[name].shape)
-                expanded_state_dict_weight[slices] = state_dict[name]
-                state_dict[name] = expanded_state_dict_weight
-
-    return state_dict
-
-
-class WatermarkEmbedder:
-    def __init__(self, watermark):
-        self.watermark = watermark
-        self.num_bits = len(WATERMARK_BITS)
-        self.encoder = WatermarkEncoder()
-        self.encoder.set_watermark("bits", self.watermark)
-
-    def __call__(self, image: torch.Tensor) -> torch.Tensor:
-        """
-        Adds a predefined watermark to the input image
-
-        Args:
-            image: ([N,] B, RGB, H, W) in range [-1, 1]
-
-        Returns:
-            same as input but watermarked
-        """
-        image = 0.5 * image + 0.5
-        squeeze = len(image.shape) == 4
-        if squeeze:
-            image = image[None, ...]
-        n = image.shape[0]
-        image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
-        # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
-        # watermarking libary expects input as cv2 BGR format
-        for k in range(image_np.shape[0]):
-            image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
-        image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
-            image.device
-        )
-        image = torch.clamp(image / 255, min=0.0, max=1.0)
-        if squeeze:
-            image = image[0]
-        image = 2 * image - 1
-        return image
-
-
-# A fixed 48-bit message that was chosen at random
-WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
-# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
-WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
-embed_watermark = WatermarkEmbedder(WATERMARK_BITS)