app.py

import os
import subprocess
import sys
import io
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import Flux2KleinPipeline
import requests
from PIL import Image
import json
import base64
from huggingface_hub import InferenceClient

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

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

hf_client = InferenceClient(
    api_key=os.environ.get("HF_TOKEN"),
)
VLM_MODEL = "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"

SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.

Guidelines:
1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
2. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.
3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.

Output only the revised prompt and nothing else."""

SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.2 by Black Forest Labs, an image-editing expert. You convert editing requests into one concise instruction (50-80 words, ~30 for brief requests).

Rules:
- Single instruction only, no commentary
- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
- Specify what changes AND what stays the same (face, lighting, composition)
- Reference actual image elements
- Turn negatives into positives ("don't change X" → "keep X")
- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
- Keep content PG-13

Output only the final instruction in plain text and nothing else."""

# Model repository IDs for 9B
REPO_ID_REGULAR = "black-forest-labs/FLUX.2-klein-base-9B"
REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"

# Load both 9B models
print("Loading 9B Regular model...")
pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)
pipe_regular.to("cuda")

print("Loading 9B Distilled model...")
pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
pipe_distilled.to("cuda")

# Dictionary for easy access
pipes = {
    "Distilled (4 steps)": pipe_distilled,
    "Base (50 steps)": pipe_regular,
}

# Default steps for each mode
DEFAULT_STEPS = {
    "Distilled (4 steps)": 4,
    "Base (50 steps)": 50,
}

DEFAULT_CFG = {
    "Distilled (4 steps)": 1.0,
    "Base (50 steps)": 4.0,
}

def image_to_data_uri(img):
    """
    Convert a PIL Image to a base64 data URI.

    Args:
        img: The PIL Image to convert.

    Returns:
        str: A data URI string containing the base64-encoded PNG image.
    """
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{img_str}"


def upsample_prompt_logic(prompt, image_list):
    """
    Enhance a text prompt using a Vision-Language Model.

    Args:
        prompt (str): The original text prompt to enhance.
        image_list: Optional list of PIL Images for context-aware enhancement.

    Returns:
        str: The enhanced prompt, or the original prompt if enhancement fails.
    """
    try:
        if image_list and len(image_list) > 0:
            # Image + Text Editing Mode
            system_content = SYSTEM_PROMPT_WITH_IMAGES
            
            # Construct user message with text and images
            user_content = [{"type": "text", "text": prompt}]
            
            for img in image_list:
                data_uri = image_to_data_uri(img)
                user_content.append({
                    "type": "image_url",
                    "image_url": {"url": data_uri}
                })
                
            messages = [
                {"role": "system", "content": system_content},
                {"role": "user", "content": user_content}
            ]
        else:
            # Text Only Mode
            system_content = SYSTEM_PROMPT_TEXT_ONLY
            messages = [
                {"role": "system", "content": system_content},
                {"role": "user", "content": prompt}
            ]

        completion = hf_client.chat.completions.create(
            model=VLM_MODEL,
            messages=messages,
            max_tokens=1024
        )
        
        return completion.choices[0].message.content
    except Exception as e:
        print(f"Upsampling failed: {e}")
        return prompt


def update_dimensions_from_image(image_list):
    """
    Update width/height based on uploaded image aspect ratio.
    
    Keeps one side at 1024 and scales the other proportionally,
    with both sides as multiples of 8.

    Args:
        image_list: Gallery list of tuples (image, caption) from Gradio.

    Returns:
        tuple: A tuple of (width, height) integers, both multiples of 8.
    """
    if image_list is None or len(image_list) == 0:
        return 1024, 1024  # Default dimensions
    
    # Get the first image to determine dimensions
    img = image_list[0][0]  # Gallery returns list of tuples (image, caption)
    img_width, img_height = img.size
    
    aspect_ratio = img_width / img_height
    
    if aspect_ratio >= 1:  # Landscape or square
        new_width = 1024
        new_height = int(1024 / aspect_ratio)
    else:  # Portrait
        new_height = 1024
        new_width = int(1024 * aspect_ratio)
    
    # Round to nearest multiple of 8
    new_width = round(new_width / 8) * 8
    new_height = round(new_height / 8) * 8
    
    # Ensure within valid range (minimum 256, maximum 1024)
    new_width = max(256, min(1024, new_width))
    new_height = max(256, min(1024, new_height))
    
    return new_width, new_height


def update_steps_from_mode(mode_choice):
    """
    Update inference steps and guidance scale based on the selected mode.

    Args:
        mode_choice (str): The selected mode, either "Distilled (4 steps)" or "Base (50 steps)".

    Returns:
        tuple: A tuple of (num_inference_steps, guidance_scale).
    """
    return DEFAULT_STEPS[mode_choice], DEFAULT_CFG[mode_choice]


@spaces.GPU(duration=85)
def infer(
    prompt: str,
    input_images=None,
    mode_choice: str = "Distilled (4 steps)",
    seed: int = 42,
    randomize_seed: bool = False,
    width: int = 1024,
    height: int = 1024,
    num_inference_steps: int = 4,
    guidance_scale: float = 4.0,
    prompt_upsampling: bool = False,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generate or edit images using FLUX.2 Klein 9B model.
    
    This tool can generate images from text prompts, or edit/combine existing images
    based on text instructions. Use the distilled mode for fast 4-step generation,
    or base mode for higher quality 50-step generation.

    Args:
        prompt (str): Text description of the image to generate, or editing instructions when input images are provided.
        input_images: Optional list of input images for editing or combining. Provide image URLs.
        mode_choice (str): Model mode - "Distilled (4 steps)" for fast generation or "Base (50 steps)" for higher quality.
        seed (str): Random seed for reproducible generation. Use "0" with randomize_seed=True for random results.
        randomize_seed (str): Set to "true" to use a random seed, "false" to use the specified seed.
        width (str): Output image width in pixels (256-1024, must be multiple of 8).
        height (str): Output image height in pixels (256-1024, must be multiple of 8).
        num_inference_steps (str): Number of denoising steps. Use "4" for distilled mode, "50" for base mode.
        guidance_scale (str): How closely to follow the prompt. Use "1.0" for distilled, "4.0" for base mode.
        prompt_upsampling (str): Set to "true" to automatically enhance the prompt using a VLM.

    Returns:
        tuple: A tuple containing the generated PIL Image and the seed used.
    """
    # Convert string inputs to proper types for MCP compatibility
    if isinstance(seed, str):
        seed = int(seed)
    if isinstance(randomize_seed, str):
        randomize_seed = randomize_seed.lower() == "true"
    if isinstance(width, str):
        width = int(width)
    if isinstance(height, str):
        height = int(height)
    if isinstance(num_inference_steps, str):
        num_inference_steps = int(num_inference_steps)
    if isinstance(guidance_scale, str):
        guidance_scale = float(guidance_scale)
    if isinstance(prompt_upsampling, str):
        prompt_upsampling = prompt_upsampling.lower() == "true"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Select the appropriate pipeline based on mode choice
    pipe = pipes[mode_choice]
    
    # Prepare image list (convert None or empty gallery to None)
    image_list = None
    if input_images is not None and len(input_images) > 0:
        image_list = []
        for item in input_images:
            image_list.append(item[0])

    # 1. Upsampling (Network bound)
    final_prompt = prompt
    if prompt_upsampling:
        progress(0.1, desc="Upsampling prompt...")
        final_prompt = upsample_prompt_logic(prompt, image_list)
        print(f"Original Prompt: {prompt}")
        print(f"Upsampled Prompt: {final_prompt}")

    # 2. Image Generation
    progress(0.2, desc=f"Generating image with 9B {mode_choice}...")
    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    pipe_kwargs = {
        "prompt": final_prompt,
        "height": height,
        "width": width,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
    }
    
    # Add images if provided
    if image_list is not None:
        pipe_kwargs["image"] = image_list
    
    image = pipe(**pipe_kwargs).images[0]
    
    return image, seed


examples = [
    ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
    ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
]

examples_images = [
    ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 1200px;
}
.gallery-container img{
    object-fit: contain;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.2 [Klein] - 9B
FLUX.2 [Klein] is a distilled model capable of generating, editing and combining images based on text instructions [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B)], [[blog](https://bfl.ai/blog/flux-2)]
        """)
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=2,
                        placeholder="Enter your prompt",
                        container=False,
                        scale=3
                    )
                    
                    run_button = gr.Button("Run", scale=1)
                    
                with gr.Accordion("Input image(s) (optional)", open=False):
                    input_images = gr.Gallery(
                        label="Input Image(s)",
                        type="pil",
                        columns=3,
                        rows=1,
                    )

                mode_choice = gr.Radio(
                    label="Mode",
                    choices=["Distilled (4 steps)", "Base (50 steps)"],
                    value="Distilled (4 steps)",
                )
                
                with gr.Accordion("Advanced Settings", open=False):
                    
                    prompt_upsampling = gr.Checkbox(
                        label="Prompt Upsampling",
                        value=False,
                        info="Automatically enhance the prompt using a VLM"
                    )
        
                    seed = gr.Slider(
                        label="Seed",
                        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=8,
                            value=1024,
                        )
                        
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=8,
                            value=1024,
                        )
                    
                    with gr.Row():
                        
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=4,
                        )
                        
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.0,
                        )
                
                
            with gr.Column():
                result = gr.Image(label="Result", show_label=False)
            
        
        gr.Examples(
            examples=examples,
            fn=infer,
            inputs=[prompt],
            outputs=[result, seed],
            cache_examples=True,
            cache_mode="lazy"
        )

        gr.Examples(
            examples=examples_images,
            fn=infer,
            inputs=[prompt, input_images],
            outputs=[result, seed],
            cache_examples=True,
            cache_mode="lazy"
        )

    # Auto-update dimensions when images are uploaded
    input_images.upload(
        fn=update_dimensions_from_image,
        inputs=[input_images],
        outputs=[width, height]
    )
    
    # Auto-update steps when mode changes
    mode_choice.change(
        fn=update_steps_from_mode,
        inputs=[mode_choice],
        outputs=[num_inference_steps, guidance_scale]
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
        outputs=[result, seed],
        api_name="generate"  # Explicit API name for MCP tool
    )

# Launch with MCP server enabled
demo.launch(mcp_server=True)