Blame for fibo2_edit_turbo_t2i.py

#!/usr/bin/env python3
"""Text-to-image with Bria's FIBO 2.0 Edit Turbo, using diffusers.

Bria distils Fibo-2.0-edit-turbo for 4-step image editing. Called with no input image, the same pipeline also generates
images from text alone: 4 steps, no classifier-free guidance.

FIBO 2.0 is trained on structured JSON captions (Bria's Visual GenAI Language, VGL). By default this script writes that
caption for you with the pipeline's own text encoder, Qwen3-VL-4B-Instruct, given the open "generate" instructions of
Bria's bria-vgl package (https://github.com/Bria-AI/vgl, MIT licence, fetched once and cached). No API key is needed.
The caption is then rendered, and saved next to the image so the result can be reproduced or edited.

Examples:
    python fibo2_edit_turbo_t2i.py "a red vintage bicycle leaning against a whitewashed wall on a narrow street"
    python fibo2_edit_turbo_t2i.py "an origami whale above a paper city" --width 1424 --height 800 --seed 42
    python fibo2_edit_turbo_t2i.py --json assets/lake-pier.json --seed 4242     # a finished caption: no writing step
    python fibo2_edit_turbo_t2i.py --plain "A lighthouse on a cliff at dusk."      # the prompt as typed, no caption
    python fibo2_edit_turbo_t2i.py --sample 2 --seed 7 "a lighthouse on a cliff at dusk"   # another take on the caption
    python fibo2_edit_turbo_t2i.py --prompts examples/prompts.json --out-dir examples --seed 4242   # a batch

Requirements: torch, diffusers with FIBO 2.0 (git main from 20 Oct 2026), transformers with Qwen3-VL (4.57 or newer),
accelerate, huggingface_hub, Pillow 10.1 or newer. The model is gated: accept its licence at
https://huggingface.co/briaai/Fibo-2.0-edit-turbo and log in (`hf auth login`) before the first run.
"""
import argparse
import json
import os
import random
import re
import sys
import time
import urllib.request

MODEL_ID = "briaai/Fibo-2.0-edit-turbo"
#: bria-vgl's "generate" instructions (the "thales" dialect): the system prompt, its field guide and the user turn
VGL_URL = "https://raw.githubusercontent.com/Bria-AI/vgl/dev/src/bria_vgl/dialects/thales"
VGL_FILES = {"system": "prompts/generate/system_prompt.txt", "fields": "input_schema/generate.txt",
             "user": "prompts/generate/final_prompt.txt"}
FONT_NAMES = ("Aileron-Regular.otf", "DejaVuSans.ttf")


def vgl_prompts(cache_dir):
    """(system prompt, user template): bria-vgl's generate prompts, downloaded once into cache_dir."""
    os.makedirs(cache_dir, exist_ok=True)
    text = {}
    for key, rel in VGL_FILES.items():
        path = os.path.join(cache_dir, rel.replace("/", "__"))
        if not os.path.isfile(path):
            with urllib.request.urlopen(f"{VGL_URL}/{rel}", timeout=60) as r:
                data = r.read().decode("utf-8")
            with open(path, "w", encoding="utf-8") as f:
                f.write(data)
        with open(path, encoding="utf-8") as f:
            text[key] = f.read()
    system = text["system"].replace("{style_preset_block}", "").replace("{llm_input_schema}", text["fields"])
    return system, text["user"]


def extract_json(reply):
    """The first JSON object in a model reply (code fences and <think> blocks removed), or None."""
    t = re.sub(r"<think>.*?</think>", "", reply or "", flags=re.S).strip()
    m = re.search(r"```(?:json)?\s*(.*?)```", t, flags=re.S)
    if m:
        t = m.group(1).strip()
    start = t.find("{")
    if start < 0:
        return None
    depth, in_str, esc = 0, False, False
    for i in range(start, len(t)):
        ch = t[i]
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
        elif ch == '"':
            in_str = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                try:
                    obj = json.loads(t[start:i + 1])
                except ValueError:
                    return None
                return obj if isinstance(obj, dict) else None
    return None


def write_caption(pipe, prompt, system, user_template, max_new_tokens=2048, sample_seed=None):
    """The VGL caption of `prompt`, written by the pipeline's own Qwen3-VL-4B-Instruct, or None: greedily (the same
    prompt, the same caption), or sampled at temperature 0.8 / top-p 0.9 from `sample_seed` (another take)."""
    import torch
    user = user_template.replace("{user_prompt}", " ".join(prompt.split()))
    messages = [{"role": "system", "content": system}, {"role": "user", "content": [{"type": "text", "text": user}]}]
    text = pipe.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = pipe.processor(text=[text], return_tensors="pt").to(pipe._execution_device)
    gen = dict(max_new_tokens=max_new_tokens, do_sample=sample_seed is not None)
    if sample_seed is not None:
        torch.manual_seed(sample_seed)
        gen.update(temperature=0.8, top_p=0.9)
    with torch.no_grad():
        out = pipe.text_encoder.generate(**inputs, **gen)
    reply = pipe.processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
    return extract_json(reply)


def font_paths(model):
    """The fonts FIBO 2.0 draws caption text with (for text_render entries), from the model folder or the Hub."""
    if os.path.isdir(model):
        return [p for p in (os.path.join(model, "fonts", n) for n in FONT_NAMES) if os.path.isfile(p)]
    try:
        from huggingface_hub import hf_hub_download
        return [hf_hub_download(model, n, subfolder="fonts") for n in FONT_NAMES]
    except Exception:
        return []


def pipeline_prompt(pipe, text):
    """The prompt as BriaFibo2Pipeline encodes it: a caption cleaned and serialised compactly, any other text as it is."""
    caption, _ = pipe._parse_prompt(text)
    return text if isinstance(caption, str) else json.dumps(caption or {}, separators=(",", ":"), ensure_ascii=False)


def draws_text(pipe, text):
    """Whether the caption asks for text: the pipeline then draws it as images and needs the prompt itself to place them."""
    caption, _ = pipe._parse_prompt(text)
    return isinstance(caption, dict) and any(isinstance(e, dict) and isinstance(e.get("text"), str) and e["text"]
                                             for e in (caption.get("text_render") or []))


def free_memory(device):
    import gc

    import torch
    gc.collect()
    if str(device).startswith("mps"):
        torch.mps.empty_cache()
    elif str(device).startswith("cuda"):
        torch.cuda.empty_cache()


def decode(pipe, latents):
    """A final latent -> PIL, as BriaFibo2Pipeline decodes it: the VAE's batch-norm statistics undone on the packed latent,
    then the VAE."""
    import torch
    with torch.no_grad():
        packed = pipe._patchify_latents(latents)
        mean = pipe.vae.bn.running_mean.view(1, -1, 1, 1).to(packed.device, packed.dtype)
        scale = (1 / torch.sqrt(pipe.vae.bn.running_var + pipe.vae.config.batch_norm_eps)).view(1, -1, 1, 1)
        packed = packed / scale.to(packed.device, packed.dtype) + mean
        image = pipe.vae.decode(pipe._unpatchify_latents(packed).to(pipe.vae.dtype), return_dict=False)[0]
    return pipe.image_processor.postprocess(image, output_type="pil")[0]


def pick_device(name):
    import torch
    if name != "auto":
        return name
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def read_batch(path):
    """[(name, prompt)] from a JSON list (strings, or objects with 'prompt' and an optional 'slug' or 'name') or a text
    file with one prompt per line (blank lines and lines starting with # skipped)."""
    with open(path, encoding="utf-8") as f:
        text = f.read()
    try:
        items = json.loads(text)
    except ValueError:
        items = [line.strip() for line in text.splitlines() if line.strip() and not line.lstrip().startswith("#")]
    jobs = []
    for i, item in enumerate(items, 1):
        if isinstance(item, dict):
            jobs.append((str(item.get("slug") or item.get("name") or f"{i:02d}"), item["prompt"]))
        else:
            jobs.append((f"{i:02d}", str(item)))
    return jobs


def main():
    ap = argparse.ArgumentParser(description="Text-to-image with FIBO 2.0 Edit Turbo (diffusers)")
    ap.add_argument("prompt", nargs="?", help="what to generate (a plain description)")
    ap.add_argument("--json", help="a finished VGL caption: a JSON file holding the caption object, or an object with "
                                   "a 'caption' key (like assets/lake-pier.json)")
    ap.add_argument("--prompts", help="a batch: a JSON list of prompts (strings, or objects with 'prompt' and an "
                                      "optional 'slug' used as the file name) or a text file with one prompt per line")
    ap.add_argument("--out-dir", default=".", help="where a batch's images and captions are saved (default: here)")
    ap.add_argument("--plain", action="store_true", help="send the prompt as typed instead of writing a VGL caption")
    ap.add_argument("--sample", nargs="?", type=int, const=-1, default=None, metavar="CAPTION_SEED",
                    help="write the caption by sampling (temperature 0.8, top-p 0.9) instead of greedily: another take "
                         "on the same prompt, from CAPTION_SEED (default: the image's seed); saved in the JSON")
    ap.add_argument("--model", default=MODEL_ID, help=f"Hub id or local diffusers folder (default {MODEL_ID})")
    ap.add_argument("--width", type=int, default=1024, help="multiple of 16, about 1 megapixel in total")
    ap.add_argument("--height", type=int, default=1024, help="multiple of 16, about 1 megapixel in total")
    ap.add_argument("--steps", type=int, default=4)
    ap.add_argument("--guidance", type=float, default=1.0, help="1.0 = no classifier-free guidance (Turbo's setting)")
    ap.add_argument("--seed", type=int, default=None, help="random for each image when omitted")
    ap.add_argument("--out", default=None, help="output PNG of a single run (default fibo2_edit_turbo_t2i_<seed>.png)")
    ap.add_argument("--device", default="auto", help="auto, cuda, mps or cpu")
    ap.add_argument("--offload", action="store_true",
                    help="model CPU offload: one component on the GPU at a time (less GPU memory, slower)")
    ap.add_argument("--vae-tiling", action="store_true",
                    help="decode the image in tiles: a lower memory peak at the end of each image")
    ap.add_argument("--cache", default=os.path.join(os.path.expanduser("~"), ".cache", "fibo2-vgl"),
                    help="where bria-vgl's prompt files are cached")
    a = ap.parse_args()
    if not (a.prompt or a.json or a.prompts):
        ap.error("give a prompt, --json with a VGL caption, or --prompts with a batch")

    import torch
    from diffusers import BriaFibo2Pipeline
    from diffusers.pipelines.bria_fibo2.pipeline_bria_fibo2 import DEFAULT_NEGATIVE_PROMPT

    width, height = a.width // 16 * 16, a.height // 16 * 16
    if (width, height) != (a.width, a.height):
        print(f"size rounded down to multiples of 16: {width}x{height}")

    # the jobs: (output PNG, prompt, caption or None)
    if a.prompts:
        os.makedirs(a.out_dir, exist_ok=True)
        jobs = [(os.path.join(a.out_dir, name + ".png"), prompt, None) for name, prompt in read_batch(a.prompts)]
    elif a.json:
        with open(a.json, encoding="utf-8") as f:
            data = json.load(f)
        caption = data["caption"] if isinstance(data.get("caption"), dict) else data
        jobs = [(a.out, data.get("prompt") if isinstance(data.get("caption"), dict) else None, caption)]
    else:
        jobs = [(a.out, a.prompt, None)]
    todo = [j for j in jobs if not (j[0] and os.path.isfile(j[0]))]
    if len(todo) < len(jobs):
        print(f"{len(jobs) - len(todo)} image(s) already there, skipped")
    if not todo:
        return 0

    device = pick_device(a.device)
    t0 = time.perf_counter()
    pipe = BriaFibo2Pipeline.from_pretrained(a.model, torch_dtype=torch.bfloat16)
    if a.offload:
        pipe.enable_model_cpu_offload(device=device)
    else:
        pipe.to(device)
    if a.vae_tiling:
        pipe.vae.enable_tiling()
    print(f"loaded {a.model} on {device} in {time.perf_counter() - t0:.0f} s")
    fonts = font_paths(a.model) or None
    system = user_template = None

    # One stage at a time, as far as the pipeline allows: 1. the text encoder writes the captions and embeds the prompts,
    # then leaves; 2. the transformer denoises every image to its latent, then leaves; 3. the VAE decodes them. A caption
    # that asks for text keeps the encoder through 2 (the pipeline shows its text images from noise 0.9 down, which needs
    # the prompt itself); --offload has the pipeline move the parts itself.
    runs = []
    for k, (out, prompt, caption) in enumerate(todo, 1):
        tag = f"[{k}/{len(todo)}] " if len(todo) > 1 else ""
        seed = a.seed if a.seed is not None else random.randrange(2 ** 31)
        caption_seed = None
        if caption is None and not a.plain:
            if system is None:
                system, user_template = vgl_prompts(a.cache)
            if a.sample is not None:
                caption_seed = seed if a.sample < 0 else a.sample
            t0 = time.perf_counter()
            caption = write_caption(pipe, prompt, system, user_template, sample_seed=caption_seed)
            if caption is None:
                print(f"{tag}no usable caption was written: the prompt is used as typed")
            else:
                print(f"{tag}caption written in {time.perf_counter() - t0:.0f} s: "
                      f"{caption.get('short_description', '')[:120]}")
        text = json.dumps(caption, separators=(",", ":"), ensure_ascii=False) if caption else prompt
        runs.append({"out": out, "prompt": prompt, "caption": caption, "caption_seed": caption_seed, "seed": seed,
                     "text": text})
    staged = not a.offload and not any(draws_text(pipe, r["text"]) for r in runs)
    negative = None
    if staged:
        with torch.no_grad():
            for r in runs:
                r["embeds"] = pipe.encode_prompt(pipeline_prompt(pipe, r["text"]), device=device)
            if a.guidance > 1.0:
                negative = pipe.encode_prompt(DEFAULT_NEGATIVE_PROMPT, device=device)
        pipe.text_encoder = None
        free_memory(device)

    finals = []
    for k, r in enumerate(runs, 1):
        tag = f"[{k}/{len(runs)}] " if len(runs) > 1 else ""
        t0 = time.perf_counter()
        kw = dict(num_inference_steps=a.steps, guidance_scale=a.guidance, width=width, height=height,
                  generator=torch.Generator("cpu").manual_seed(r["seed"]), output_type="latent")
        if staged:
            kw.update(prompt_embeds=r["embeds"][0], prompt_embeds_mask=r["embeds"][1])
            if negative is not None:
                kw.update(negative_prompt_embeds=negative[0], negative_prompt_embeds_mask=negative[1])
        else:
            kw.update(prompt=r["text"], fonts=fonts)
        finals.append(pipe(**kw).images)
        r.pop("embeds", None)
        print(f"{tag}denoised in {time.perf_counter() - t0:.0f} s")
    if not a.offload:
        pipe.transformer = pipe.text_encoder = None
        free_memory(device)

    for k, (r, latents) in enumerate(zip(runs, finals), 1):
        tag = f"[{k}/{len(runs)}] " if len(runs) > 1 else ""
        t0 = time.perf_counter()
        image = decode(pipe, latents)
        out = r["out"] or f"fibo2_edit_turbo_t2i_{r['seed']}.png"
        image.save(out)
        record = {"model": a.model, "seed": r["seed"], "steps": a.steps, "guidance_scale": a.guidance, "width": width,
                  "height": height, "prompt": r["prompt"], "caption": r["caption"]}
        if r["caption_seed"] is not None and r["caption"]:
            record["caption_seed"] = r["caption_seed"]    # written with --sample: reproduce with --json, or this seed
        with open(os.path.splitext(out)[0] + ".json", "w", encoding="utf-8") as f:
            json.dump(record, f, indent=2, ensure_ascii=False)
        print(f"{tag}saved {out} ({width}x{height}, {a.steps} steps, seed {r['seed']}; decoded in "
              f"{time.perf_counter() - t0:.0f} s)")
    return 0


if __name__ == "__main__":
    sys.exit(main())