divisor.flux2.prompt

  1# SPDX-License-Identifier: MPL-2.0 AND LicenseRef-Commons-Clause-License-Condition-1.0
  2# <!-- // /*  d a r k s h a p e s */ -->
  3# adapted BFL Flux code from https://github.com/black-forest-labs/flux2
  4
  5from dataclasses import replace
  6
  7import torch
  8from fire import Fire
  9from nnll.console import nfo
 10from PIL import Image
 11
 12from divisor.controller import rng
 13from divisor.flux1.loading import load_ae, load_flow_model, load_mistral_small_embedder
 14from divisor.flux1.prompt import parse_prompt
 15from divisor.flux2.sampling import (
 16    batched_prc_img,
 17    batched_prc_txt,
 18    denoise_interactive,
 19    encode_image_refs,
 20    get_schedule,
 21)
 22from divisor.noise import get_noise
 23from divisor.registry import gfx_device, gfx_dtype, empty_cache
 24from divisor.spec import ModelSpec, flux_configs, get_model_spec
 25from divisor.state import InferenceState, MenuState
 26
 27
 28def main(
 29    mir_id: str = "model.dit.flux2-dev",
 30    ae_id: str = "model.vae.flux2-dev",
 31    width: int = 1360,
 32    height: int = 768,
 33    guidance: float = 4,
 34    seed: int = rng.next_seed(),
 35    prompt: str = "",
 36    quantization: bool = False,
 37    device: torch.device = gfx_device,
 38    num_steps: int = 50,
 39    upsample_prompt: bool = False,
 40    loop: bool = False,
 41    offload: bool = False,
 42    compile: bool = False,
 43    verbose: bool = False,
 44    input_images: list[str] | None = None,
 45) -> None:
 46    """Sample the flux model. Either interactively (set `--loop`) or run for a single image.\n
 47    :param name: Name of the model to load
 48    :param height: height of the sample in pixels (should be a multiple of 16)
 49    :param width: width of the sample in pixels (should be a multiple of 16)
 50    :param seed: Set a seed for sampling
 51    :param output_name: where to save the output image, `{idx}` will be replaced by the index of the sample
 52    :param prompt: Prompt used for sampling
 53    :param device: Pytorch device
 54    :param num_steps: number of sampling steps (default 4 for schnell, 28 for guidance distilled)
 55    :param loop: start an interactive session and sample multiple times
 56    :param guidance: guidance value used for guidance distillation
 57    """
 58
 59    precision = gfx_dtype
 60    prompt_parts = prompt.split("|")
 61    if len(prompt_parts) == 1:
 62        prompt = prompt_parts[0]
 63        additional_prompts = None
 64    else:
 65        additional_prompts = prompt_parts[1:]
 66        prompt = prompt_parts[0]
 67
 68    mistral = load_mistral_small_embedder()
 69    if quantization:
 70        mir_id += ":*@fp8-sai"
 71    model_spec: ModelSpec = get_model_spec(mir_id, flux_configs)
 72    ae_spec = get_model_spec(ae_id, flux_configs)
 73    model = load_flow_model(
 74        model_spec,
 75        device=torch.device("cpu") if offload else device,
 76        verbose=verbose,
 77    )
 78
 79    is_compiled = False
 80    if compile and not offload:  # Compiled models can't be easily moved between devices
 81        nfo("Compilation enabled.")
 82        model = torch.compile(model)  # type: ignore[assignment]
 83        is_compiled = True
 84
 85    ae = load_ae(ae_spec, device=torch.device("cpu") if offload else device)
 86    ae.eval()
 87    mistral.eval()
 88
 89    state = MenuState.from_cli_args(
 90        prompt=prompt,
 91        width=width,
 92        height=height,
 93        num_steps=num_steps,
 94        guidance=guidance,
 95        seed=seed,
 96    )
 97
 98    if loop:
 99        state = parse_prompt(state)
100
101    while state is not None:
102        if state.seed is None:
103            seed = rng.next_seed()
104            state = replace(state, seed=seed)
105        else:
106            rng.next_seed(state.seed)
107        # At this point, state.seed is guaranteed to be an int
108        assert state.seed is not None, "Seed must be set"
109        assert state.width is not None and state.height is not None, "Width and height must be set"
110        assert state.num_steps is not None, "num_steps must be set"
111        assert state.prompt is not None, "Prompt must be set"
112        nfo(f"Generating with seed {rng.seed}: {state.prompt}")
113
114        x = get_noise(
115            1,
116            state.height,
117            state.width,
118            dtype=precision,
119            seed=rng.seed,  # type: ignore
120            device=device,
121            version_2=True,
122        )
123
124        with torch.no_grad():
125            if input_images:
126                img_ctx = [Image.open(input_image) for input_image in input_images]
127                ref_tokens, ref_ids = encode_image_refs(ae, img_ctx)  # type: ignore
128            else:
129                ref_tokens = None
130                ref_ids = None
131
132            if upsample_prompt:
133                # Use local model for upsampling
134                upsampled_prompts = mistral.upsample_prompt([state.prompt], img=[img_ctx] if img_ctx else None)  # type: ignore
135                prompt = upsampled_prompts[0] if upsampled_prompts else state.prompt
136                state = replace(state, prompt=prompt)
137            else:
138                prompt = state.prompt
139
140            ctx = mistral([prompt]).to(precision)  # type: ignore
141            ctx, ctx_ids = batched_prc_txt(ctx)
142
143            randn = get_noise(
144                1,
145                state.height,  # type: ignore
146                state.width,  # type: ignore
147                dtype=torch.bfloat16,
148                seed=state.seed,  # type: ignore
149                device=device,
150                version_2=True,
151            )
152
153            x, x_ids = batched_prc_img(randn)
154
155            timesteps = get_schedule(state.num_steps, x.shape[1])  # type: ignore
156            # Update state with runtime information
157            state = state.with_runtime_state(
158                current_timestep=0.0,
159                current_sample=x,
160                timestep_index=0,
161                total_timesteps=len(timesteps),
162            )
163
164            if offload:
165                if device.type == "cuda":
166                    mistral = mistral.cpu()  # type: ignore
167                else:
168                    mistral = None
169                    del mistral
170                empty_cache
171                if is_compiled:
172                    raise RuntimeError("Cannot use offload=True with compile=True. Compile after model is on device, or disable compilation when offloading.")
173                model = model.to(device)  # type: ignore[attr-defined]
174                if compile:
175                    nfo("Compiling model on device.")
176                    model = torch.compile(model, mode="max-autotune")  # type: ignore[assignment]
177                    is_compiled = True
178
179            x = denoise_interactive(
180                model,  # type: ignore
181                InferenceState(
182                    img=x,
183                    img_ids=x_ids,
184                    txt=ctx,
185                    txt_ids=ctx_ids,
186                    timesteps=timesteps,
187                    img_cond_seq=ref_tokens,
188                    img_cond_seq_ids=ref_ids,
189                    state=state,
190                    ae=ae,
191                ),
192            )
193            if loop:
194                nfo("-" * 80)
195                state = parse_prompt(state)
196            elif additional_prompts:
197                next_prompt = additional_prompts.pop(0)
198                state = replace(state, prompt=next_prompt)
199            else:
200                state = None
201
202
203if __name__ == "__main__":
204    Fire(main)