import json

import torch
from diffusers import BriaFibo2Pipeline

pipe = BriaFibo2Pipeline.from_pretrained("briaai/Fibo-2.0-dev", dtype=torch.bfloat16).to("cuda")

caption = {
    "short_description": "A red vintage bicycle leaning against a whitewashed wall on a narrow Mediterranean street.",
    "objects": [
        {
            "description": "A red vintage bicycle with a wicker basket of lemons on its handlebars.",
            "location": "center",
            "relative_size": "large within frame",
        }
    ],
    "background_setting": "A whitewashed wall with a blue wooden door and potted red geraniums.",
    "lighting": {"conditions": "bright morning sun", "direction": "from the right", "shadows": "long and crisp"},
    "aesthetics": {"composition": "centered", "color_scheme": "red, white and blue", "mood_atmosphere": "calm"},
    "photographic_characteristics": {"depth_of_field": "deep", "focus": "sharp on the bicycle", "camera_angle": "eye level"},
    "style_medium": "photograph",
    "context": "A travel photograph.",
    "artistic_style": "realistic",
}
image = pipe(json.dumps(caption)).images[0]
image.save("bicycle.png")

# The caption's short description alone, as plain text
image = pipe("A red vintage bicycle leaning against a whitewashed wall on a narrow Mediterranean street.").images[0]

# Exactly a "short_description" and a "text_render" list. The list is kept even when it's empty, to say the image
# has no text
caption = {"short_description": "A bakery storefront at dawn with a painted sign.", "text_render": [{"text": "FRESH BREAD"}]}
image = pipe(json.dumps(caption)).images[0]

import torch
from diffusers import BriaFibo2Pipeline
from diffusers.utils import load_image

pipe = BriaFibo2Pipeline.from_pretrained("briaai/Fibo-2.0-dev", dtype=torch.bfloat16).to("cuda")
image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
)
edited = pipe("Make the background a snowy forest", image=image).images[0]
edited.save("pikachu_snow.png")

import torch
from diffusers import BriaFibo2Pipeline
from diffusers.utils import load_image

pipe = BriaFibo2Pipeline.from_pretrained("briaai/Fibo-2.0-edit-turbo", dtype=torch.bfloat16).to("cuda")
image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
)
edited = pipe("Make the background a snowy forest", image=image, num_inference_steps=4, guidance_scale=1.0).images[0]
edited.save("pikachu_snow.png")

import json

prompt = {
    "short_description": "A crocheted yellow Pikachu doll sitting in a snowy pine forest.",
    "objects": [{"description": "A crocheted yellow Pikachu doll with red cheeks.", "location": "center"}],
    "background_setting": "A snowy pine forest with falling snow.",
    "edit_instruction": "Make the background a snowy forest",
}
edited = pipe(json.dumps(prompt), image=image).images[0]

prompt = {
    "structured_caption": "A crocheted yellow Pikachu doll sitting in a snowy pine forest.",
    "edit_instruction": "Make the background a snowy forest",
}
edited = pipe(json.dumps(prompt), image=image).images[0]

prompt = {
    "short_description": "A crocheted yellow Pikachu doll holding a small wooden sign.",
    "text_render": [{"text": "HELLO"}],
    "edit_instruction": "Give the doll a small wooden sign that reads HELLO",
}
edited = pipe(json.dumps(prompt), image=image).images[0]

cat = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
edited = pipe("Put the cat from <image_2> next to the doll in <image_1>", image=[image, cat]).images[0]

edited = pipe("Make the background a snowy forest", image=image, width=1424, height=800).images[0]

from PIL import Image, ImageDraw

mask = Image.new("L", image.size, 0)
ImageDraw.Draw(mask).rectangle((image.width // 2, 0, image.width, image.height // 2), fill=255)
edited = pipe("Add a red balloon floating in the sky", image=image, mask=mask).images[0]

prompt = json.dumps(caption)
image = pipe(prompt, render_text_images=False).images[0]

from huggingface_hub import hf_hub_download

fonts = [
    hf_hub_download("briaai/Fibo-2.0-dev", name, subfolder="fonts")
    for name in ("Aileron-Regular.otf", "DejaVuSans.ttf")
]
image = pipe(prompt, fonts=fonts).images[0]

pipe.transformer.set_attention_backend("_flash_3_varlen_hub")