app.py

import torch
import spaces
import gradio as gr
import random
import numpy as np
import os

from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from huggingface_hub import hf_hub_download

# -----------------------------------------------------------------------------
# Config — edit these to match your checkpoint
# -----------------------------------------------------------------------------
HF_TOKEN = os.getenv("HF_TOKEN")

# Key-patched raw checkpoint (QK-norm keys renamed .weight -> .scale so
# diffusers' from_single_file converter can parse it -- see
# convert_and_upload.py / upload_fixed_checkpoint.py for how this was produced).
CUSTOM_CHECKPOINT_REPO = "NickSh228/Fluxy"   # <-- set after running upload_fixed_checkpoint.py
CUSTOM_CHECKPOINT_FILE = "fluxed-fixed.safetensors"

# Text encoders + VAE come from the base FLUX.1-dev repo, since these files are
# identical for essentially every Flux.1-dev based checkpoint (CLIP-L, T5-XXL, ae.safetensors).
BASE_REPO = "black-forest-labs/FLUX.1-dev"

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device}")

MAX_SEED = np.iinfo(np.int32).max

# -----------------------------------------------------------------------------
# Load pipeline
# -----------------------------------------------------------------------------
# 1) Custom transformer -- raw checkpoint with pre-patched QK-norm key names.
# Download locally first (avoids the from_single_file URL-doubling bug), then
# convert via from_single_file. This conversion step needs ~24GB+ of RAM/VRAM,
# which the Space's hardware should have even though a laptop might not.
checkpoint_path = hf_hub_download(
    repo_id=CUSTOM_CHECKPOINT_REPO,
    filename=CUSTOM_CHECKPOINT_FILE,
    token=HF_TOKEN,
)
transformer = FluxTransformer2DModel.from_single_file(
    checkpoint_path,
    config=BASE_REPO,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)

# 2) Standard Flux VAE + CLIP-L + T5-XXL text encoders, from the base repo
vae = AutoencoderKL.from_pretrained(BASE_REPO, subfolder="vae", torch_dtype=torch.bfloat16, token=HF_TOKEN)
text_encoder = CLIPTextModel.from_pretrained(BASE_REPO, subfolder="text_encoder", torch_dtype=torch.bfloat16, token=HF_TOKEN)
tokenizer = CLIPTokenizer.from_pretrained(BASE_REPO, subfolder="tokenizer", token=HF_TOKEN)
text_encoder_2 = T5EncoderModel.from_pretrained(BASE_REPO, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, token=HF_TOKEN)
tokenizer_2 = T5TokenizerFast.from_pretrained(BASE_REPO, subfolder="tokenizer_2", token=HF_TOKEN)

# 3) Scheduler — closest diffusers equivalent to ComfyUI's "beta" scheduler.
# max_shift / base_shift taken from the ModelSamplingFlux node in the reference workflow.
scheduler = FlowMatchEulerDiscreteScheduler(
    use_beta_sigmas=True,
    base_shift=0.50,
    max_shift=1.21,
)

pipe = FluxPipeline(
    scheduler=scheduler,
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    text_encoder_2=text_encoder_2,
    tokenizer_2=tokenizer_2,
    transformer=transformer,
)
pipe.to(device)

# NOTE: diffusers has no DPM++ 2M solver for Flux's flow-matching models — the
# workflow's "dpmpp_2m" sampler can't be reproduced exactly. FlowMatchEulerDiscreteScheduler
# with use_beta_sigmas=True is the closest available match and gets you most of the way there.


@spaces.GPU(duration=60)
def generate_image(prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt, progress=gr.Progress(track_tqdm=True)):
    if seed == 0:
        seed = random.randint(1, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    with torch.inference_mode():
        output = pipe(
            prompt=prompt,
            num_inference_steps=num_inference_steps,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            generator=generator,
            num_images_per_prompt=num_images_per_prompt,
        ).images

    return output


# -----------------------------------------------------------------------------
# UI
# -----------------------------------------------------------------------------
examples = [
    ["A cat holding a sign that says hello world"],
    ["a tiny astronaut hatching from an egg on the moon"],
    ["An astronaut on mars in a futuristic cyborg suit."],
]

css = '''
.gradio-container{max-width: 1000px !important}
h1{text-align:center}
'''

with gr.Blocks(css=css) as demo:
    gr.HTML("<h1>Fluxy (Fluxed) — FLUX.1-dev fine-tune</h1>")

    with gr.Group():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", info="Describe the image you want", placeholder="A cat...")
            run_button = gr.Button("Run")
        result = gr.Gallery(label="Generated AI Images", elem_id="gallery")

    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            # Defaults match the reference ComfyUI workflow (KSamplerSelect / BasicScheduler)
            num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=30, step=1)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=7.0, value=4.0, step=0.1)
        with gr.Row():
            # Defaults match the reference workflow's Empty Latent Image node
            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=32, value=768)
            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=32, value=1344)
        with gr.Row():
            seed = gr.Slider(value=42, minimum=0, maximum=MAX_SEED, step=1, label="Seed", info="0 = random")
            num_images_per_prompt = gr.Slider(label="Images Per Prompt", minimum=1, maximum=4, step=1, value=1)

    gr.Examples(
        examples=examples,
        fn=generate_image,
        inputs=[prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt],
        outputs=[result],
        cache_examples=False,
    )

    gr.on(
        triggers=[prompt.submit, run_button.click],
        fn=generate_image,
        inputs=[prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt],
        outputs=[result],
    )

demo.queue().launch(share=False)