app.py

import gradio as gr
import numpy as np
import random
import spaces
import torch
from PIL import Image
from typing import Annotated, Iterator, Tuple
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

TOOL_SUMMARY = (
    "Generate an image from a text prompt via an uncensored model; "
    "tunable model/steps/guidance/size, supports negative prompt and seed; returns a PIL.Image. "
    "Return the generated media to the user in this format `![Alt text](URL)`."
)

good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=good_vae).to(device)
pipe.load_lora_weights("enhanceaiteam/Flux-Uncensored-V2")
torch.cuda.empty_cache()

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

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

@spaces.GPU(duration=25)
def infer(
    prompt: Annotated[str, "Text description of the image to generate."],
    seed: Annotated[int, "Random seed for reproducibility. Use 0 for a random seed per call."] = 42,
    randomize_seed: Annotated[bool, "If true, pick a new random seed for every call (overrides seed)."] = True,
    width: Annotated[int, "Output width in pixels (256–2048, multiple of 32 recommended)."] = 768,
    height: Annotated[int, "Output height in pixels (256–2048, multiple of 32 recommended)."] = 768,
    guidance_scale: Annotated[float, "Classifier-free guidance scale (1–15). Higher = follow the prompt more closely."] = 4.5,
    num_inference_steps: Annotated[int, "Number of denoising steps (1–50). Higher = slower, potentially higher quality."] = 24,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> Iterator[Tuple[Image.Image, int]]:
    """
    Generate professional-quality images with Flux.1 dev paired with the Flux-Uncensored-V2 LoRA.

    Streams intermediate frames so MCP clients can preview progress. Optimized for product photography,
    fashion, concept art, and stock-style imagery where lighting and material realism matter.

    Args:
        prompt (str): Text description of the desired image.
        seed (int): Random seed for reproducibility. Use 0 to randomize per call.
        randomize_seed (bool): If true, ignore seed and pick a new random value per call.
        width (int): Image width in pixels. Upper bound enforced by MAX_IMAGE_SIZE.
        height (int): Image height in pixels. Upper bound enforced by MAX_IMAGE_SIZE.
        guidance_scale (float): Classifier-free guidance strength.
        num_inference_steps (int): Number of denoising iterations to perform.
        progress (gr.Progress): Managed by Gradio to surface progress updates.

    Yields:
        tuple[Image.Image, int]: Intermediate or final image paired with the seed used.

    Raises:
        gr.Error: If the prompt is empty or the requested dimensions exceed MAX_IMAGE_SIZE.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please provide a non-empty prompt.")
    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
        raise gr.Error(f"Width and height must be <= {MAX_IMAGE_SIZE}.")

    if randomize_seed or seed == 0:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        output_type="pil",
        good_vae=good_vae,
    ):
        yield img, seed

examples = [
    [
        "a tiny astronaut hatching from an egg on mars",
        0,
        True,
        768,
        768,
        4.5,
        24,
    ],
    [
        "a dog holding a sign that reads 'hello world'",
        0,
        True,
        768,
        768,
        4.5,
        24,
    ],
    [
        "an anime illustration of an apple strudel",
        0,
        True,
        768,
        768,
        4.5,
        24,
    ],
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 620px;
}
"""

with gr.Blocks() as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# Flux-Uncensored-V2 (LoRA)
    Flux.1 dev base model enhanced with the Flux-Uncensored-V2 LoRA for unrestricted generation.
    [[license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[base model](https://huggingface.co/black-forest-labs/FLUX.1-dev)] [[LoRA](https://huggingface.co/enhanceaiteam/Flux-Uncensored-V2)]
        """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            seed = gr.Slider(
                label="Seed (0 = random)",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
            
            with gr.Row():

                guidance_scale = gr.Slider(
                    label="CFG Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=4.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=24,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
            outputs = [result, seed],
            cache_examples="lazy"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs = [result, seed],
        api_name="generate_image",
        api_description=TOOL_SUMMARY,
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme="Nymbo/Nymbo_Theme", css=css)