app.py

import os
import subprocess
import sys
import io
import re
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
import requests
from PIL import Image
import json
import base64
from huggingface_hub import InferenceClient

subprocess.check_call([sys.executable, "-m", "pip", "install", "spaces==0.43.0"])

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."""

def remote_text_encoder(prompts):
    from gradio_client import Client
    
    client = Client("multimodalart/mistral-text-encoder")
    result = client.predict(
        prompt=prompts,
        api_name="/encode_text"
    )
    
    # Load returns a tensor, usually on CPU by default
    prompt_embeds = torch.load(result[0])
    return prompt_embeds

# Load model
repo_id = "black-forest-labs/FLUX.2-dev"

dit = Flux2Transformer2DModel.from_pretrained(
    repo_id,
    subfolder="transformer",
    torch_dtype=torch.bfloat16
)

pipe = Flux2Pipeline.from_pretrained(
    repo_id,
    text_encoder=None,
    transformer=dit,
    torch_dtype=torch.bfloat16
)
pipe.to(device)

# Pull pre-compiled Flux2 Transformer blocks from HF hub
spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")

def image_to_data_uri(img):
    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):
    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 sliders based on uploaded image aspect ratio.
    Keeps one side at 1024 and scales the other proportionally, with both sides as 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

# Updated duration function to match generate_image arguments (including progress)
def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    num_images = 0 if image_list is None else len(image_list)
    step_duration = 1 + 0.8 * num_images
    return max(65, num_inference_steps * step_duration + 10)

@spaces.GPU(duration=get_duration)
def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    # Move embeddings to GPU only when inside the GPU decorated function
    prompt_embeds = prompt_embeds.to(device)
    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    pipe_kwargs = {
        "prompt_embeds": prompt_embeds,
        "image": image_list,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
        "width": width,
        "height": height,
    }
    
    # Progress bar for the actual generation steps
    if progress:
        progress(0, desc="Starting generation...")
        
    image = pipe(**pipe_kwargs).images[0]
    return image

def parse_aspect_ratio(aspect_ratio_str):
    """Parse aspect ratio string to get width and height."""
    # Extract dimensions from format like "1:1 (1024x1024)"
    match = re.search(r'\((\d+)x(\d+)\)', aspect_ratio_str)
    if match:
        return int(match.group(1)), int(match.group(2))
    return 1024, 1024  # Default

def infer(prompt, aspect_ratio="2:3 (1344x2048)", progress=gr.Progress(track_tqdm=True)):
    """Generate an image using FLUX.2 model."""
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to generate an image.")

    # Always use random seed
    seed = random.randint(0, MAX_SEED)

    # Parse aspect ratio to get width and height
    width, height = parse_aspect_ratio(aspect_ratio)

    # Fixed inference parameters
    num_inference_steps = 30
    guidance_scale = 4.0

    # Text Encoding (Network bound - No GPU needed)
    progress(0.1, desc="Encoding prompt...")
    prompt_embeds = remote_text_encoder(prompt)

    # Image Generation (GPU bound)
    progress(0.3, desc="Generating image...")
    image = generate_image(
        prompt_embeds,
        None,  # No input images
        width,
        height,
        num_inference_steps,
        guidance_scale,
        seed,
        progress
    )

    return image

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 = [
    # ["Replace the top of the person from image 1 with the one from image 2", ["person1.webp", "woman2.webp"]],
    ["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"]]
]

# Apple-inspired CSS styling (matching Z-Image)
css = """
/* Global Styles */
.gradio-container {
    max-width: 85vw !important;
    margin: 0 auto !important;
    padding: 48px 20px !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Inter', 'Segoe UI', 'Roboto', sans-serif !important;
}

* {
    transition: none !important;
    animation: none !important;
}

/* Header */
.header-container {
    text-align: left;
    margin-bottom: 24px;
}

.main-title {
    font-size: 32px !important;
    font-weight: 600 !important;
    letter-spacing: -0.02em !important;
    line-height: 1.07 !important;
    color: #1d1d1f !important;
    margin: 0 0 16px 0 !important;
}

/* Input Section */
.input-section {
    background: #ffffff;
    border-radius: 18px;
    padding: 32px;
    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08);
}

#input-column {
    background: #ffffff !important;
    border-radius: 18px !important;
    padding: 32px !important;
    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08) !important;
    width: 550px !important;
    min-width: 550px !important;
    max-width: 550px !important;
    flex: 0 0 550px !important;
}

/* Textbox */
#prompt-textbox textarea {
    font-size: 17px !important;
    line-height: 1.47 !important;
    border-radius: 12px !important;
    border: 1px solid #d2d2d7 !important;
    padding: 12px 16px !important;
    background: #ffffff !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Inter', sans-serif !important;
    min-height: 500px !important;
    max-height: 700px !important;
    height: 500px !important;
    resize: vertical !important;
    overflow-y: auto !important;
    margin-bottom: 16px !important;
}

#prompt-textbox textarea:focus {
    border-color: #0071e3 !important;
    box-shadow: 0 0 0 4px rgba(0, 113, 227, 0.15) !important;
    outline: none !important;
}

/* Character counter */
.char-counter {
    text-align: center;
    font-size: 13px;
    color: #86868b;
    margin-top: 8px;
    margin-bottom: 16px;
}

.char-counter.warning {
    color: #ff9500;
}

.char-counter.limit {
    color: #ff3b30;
}

/* Button */
button.primary {
    font-size: 17px !important;
    font-weight: 400 !important;
    padding: 12px 32px !important;
    border-radius: 980px !important;
    background: #0071e3 !important;
    border: none !important;
    color: #ffffff !important;
    min-height: 44px !important;
    letter-spacing: -0.01em !important;
    cursor: pointer !important;
    width: 100% !important;
    margin-top: 16px !important;
}

button.primary:hover {
    background: #0077ed !important;
}

/* Output Section */
#output-column {
    background: #ffffff !important;
    border-radius: 18px !important;
    padding: 32px !important;
    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08) !important;
    overflow: hidden !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    min-height: 80vh !important;
    max-height: 90vh !important;
    flex: 1 1 auto !important;
}

#output-column img {
    border-radius: 12px !important;
    max-width: 100% !important;
    max-height: 85vh !important;
    width: auto !important;
    height: auto !important;
    object-fit: contain !important;
}

/* Dark Mode */
.dark #input-column,
.dark #output-column {
    background: #1d1d1f !important;
    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.4) !important;
}

.dark .main-title {
    color: #f5f5f7 !important;
}

.dark #prompt-textbox textarea {
    background: #1d1d1f !important;
    border-color: #424245 !important;
    color: #f5f5f7 !important;
}

.dark .char-counter {
    color: #a1a1a6 !important;
}

/* Layout - Force horizontal layout */
#main-row {
    align-items: flex-start !important;
    flex-direction: row !important;
    display: flex !important;
    flex-wrap: nowrap !important;
    width: 100% !important;
    gap: 24px !important;
}

.gradio-row {
    align-items: flex-start !important;
    flex-direction: row !important;
    display: flex !important;
    flex-wrap: nowrap !important;
    width: 100% !important;
}

.gradio-row > .gradio-column:first-child,
.gradio-row > div:first-child {
    width: 550px !important;
    min-width: 550px !important;
    max-width: 550px !important;
    flex: 0 0 550px !important;
}

.gradio-row > .gradio-column:last-child,
.gradio-row > div:last-child {
    flex: 1 1 auto !important;
    min-width: 0 !important;
}

/* Hide footer */
footer {
    display: none !important;
}

/* Responsive adjustments */
@media (max-width: 1200px) {
    #main-row {
        flex-direction: column !important;
        flex-wrap: wrap !important;
    }

    #input-column {
        width: 100% !important;
        min-width: 100% !important;
        max-width: 100% !important;
        flex: 1 1 100% !important;
    }

    #output-column {
        width: 100% !important;
        min-height: 50vh !important;
    }
}
"""

# JavaScript for layout control and character counter (matching Z-Image)
js_code = """
function() {
    function setupCharCounter() {
        const textbox = document.querySelector('#prompt-textbox textarea');
        const counter = document.querySelector('.char-counter');
        const countSpan = document.getElementById('char-count');

        if (textbox && counter && countSpan) {
            function updateCount() {
                const len = textbox.value.length;
                countSpan.textContent = len;

                counter.classList.remove('warning', 'limit');
                if (len >= 2000) {
                    counter.classList.add('limit');
                } else if (len >= 1800) {
                    counter.classList.add('warning');
                }
            }

            textbox.addEventListener('input', updateCount);
            updateCount();
        }
    }

    function forceHorizontalLayout() {
        const container = document.querySelector('.gradio-container');
        if (container) {
            container.style.maxWidth = '85vw';
            container.style.width = '85vw';
        }

        const mainRow = document.getElementById('main-row');
        if (mainRow) {
            mainRow.style.flexDirection = 'row';
            mainRow.style.flexWrap = 'nowrap';
            mainRow.style.display = 'flex';
            mainRow.style.width = '100%';
        }

        const rows = document.querySelectorAll('.gradio-row');
        rows.forEach(row => {
            row.style.flexDirection = 'row';
            row.style.flexWrap = 'nowrap';
            row.style.display = 'flex';
        });

        const inputCol = document.getElementById('input-column');
        if (inputCol) {
            inputCol.style.width = '550px';
            inputCol.style.minWidth = '550px';
            inputCol.style.maxWidth = '550px';
            inputCol.style.flex = '0 0 550px';
        }

        const outputCol = document.getElementById('output-column');
        if (outputCol) {
            outputCol.style.flex = '1 1 auto';
            outputCol.style.minWidth = '0';
        }
    }

    forceHorizontalLayout();
    setupCharCounter();

    setTimeout(forceHorizontalLayout, 100);
    setTimeout(forceHorizontalLayout, 500);
    setTimeout(forceHorizontalLayout, 1000);
    setTimeout(forceHorizontalLayout, 2000);

    setTimeout(setupCharCounter, 100);
    setTimeout(setupCharCounter, 500);

    const observer = new MutationObserver(forceHorizontalLayout);
    observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
}
"""

with gr.Blocks() as demo:

    # Inject CSS via HTML style tag
    gr.HTML(f"<style>{css}</style>")

    # Two-column layout
    with gr.Row(equal_height=False, elem_id="main-row"):

        # LEFT COLUMN - Input Controls
        with gr.Column(scale=0, min_width=550, elem_id="input-column"):

            # Header
            gr.HTML("""
            <div class="header-container">
                <h1 class="main-title">FLUX.2 [dev]</h1>
            </div>
            """)

            # Prompt Textbox
            prompt = gr.Textbox(
                placeholder="Describe the image you want to create...",
                lines=15,
                max_lines=20,
                max_length=2000,
                label="Prompt",
                show_label=True,
                container=True,
                autoscroll=False,
                elem_id="prompt-textbox",
            )

            # Character Counter
            char_counter = gr.HTML(
                '<div class="char-counter"><span id="char-count">0</span> characters (max 2000)</div>'
            )

            # Aspect Ratio Dropdown
            aspect_ratio = gr.Dropdown(
                choices=[
                    "1:1 (2048x2048)",
                    "2:3 (1344x2048)",
                    "3:2 (2048x1344)",
                    "3:4 (1536x2048)",
                    "4:3 (2048x1536)",
                    "9:16 (1152x2048)",
                    "16:9 (2048x1152)",
                ],
                value="2:3 (1344x2048)",
                label="Aspect Ratio",
                show_label=True,
                container=True,
            )

            # Generate Button
            generate_btn = gr.Button(
                "Generate",
                variant="primary",
                size="lg",
                elem_classes="primary"
            )

        # RIGHT COLUMN - Image Output
        with gr.Column(scale=2, elem_id="output-column"):
            result = gr.Image(
                label="Result",
                show_label=False,
                type="pil",
                format="png",
            )

    # Event handlers
    gr.on(
        triggers=[generate_btn.click, prompt.submit],
        fn=infer,
        inputs=[prompt, aspect_ratio],
        outputs=[result],
        show_progress="full"
    )

    # Load JavaScript for layout control
    demo.load(None, None, None, js=js_code)

demo.launch()