app.py

# -*- coding: utf-8 -*-
"""
Fluent Emoji Generator
Hugging Face Space: klibaner/fluent-emoji-generator

Gradio 6.x + ZeroGPU + SDXL LoRA
"""

from __future__ import annotations

import random
import tempfile
import time
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image


# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------

BASE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
LORA_MODEL_ID = "klibaner/fluent-emoji-style-sdxl"
LORA_WEIGHT_NAME = "microsoft-fluent-emoji-sdxl.safetensors"
LORA_ADAPTER_NAME = "fluent_emoji"
TRIGGER_PHRASE = "msemj style"

# Hard limits enforced server-side. The UI widgets below (sliders, max_length)
# only constrain well-behaved browser clients -- anyone calling the exposed
# `generate` API endpoint directly can send arbitrary values, so every one of
# these bounds is re-checked inside generate() itself.
MAX_PROMPT_LENGTH = 700
MAX_NEGATIVE_PROMPT_LENGTH = 500
MIN_STEPS, MAX_STEPS = 15, 50
MIN_GUIDANCE, MAX_GUIDANCE = 1.0, 12.0
MIN_LORA_WEIGHT, MAX_LORA_WEIGHT = 0.0, 1.5
MIN_SEED, MAX_SEED = 0, 2_147_483_647

DEFAULT_PROMPT = (
    "msemj style, a friendly blue robot holding a laptop, "
    "3d emoji icon, smooth rounded shapes, soft studio lighting, "
    "isolated on white background"
)

DEFAULT_NEGATIVE_PROMPT = (
    "text, letters, watermark, logo, blurry, low quality, distorted, "
    "deformed, cropped, duplicate objects, busy background"
)

EXAMPLES = [
    [
        "msemj style, a friendly blue robot holding a laptop, "
        "3d emoji icon, smooth rounded shapes, isolated on white background"
    ],
    [
        "msemj style, a golden trophy with stars, "
        "3d emoji icon, isolated on white background"
    ],
    [
        "msemj style, a small online store with a striped awning, "
        "3d emoji icon, isolated on white background"
    ],
    [
        "msemj style, a red rocket launching from a cloud, "
        "3d emoji icon, isolated on white background"
    ],
    [
        "msemj style, a purple artificial intelligence brain with small circuits, "
        "3d emoji icon, smooth rounded shapes, isolated on white background"
    ],
    [
        "msemj style, a cheerful orange cat wearing headphones, "
        "3d emoji icon, smooth rounded shapes, isolated on white background"
    ],
]


# ---------------------------------------------------------------------------
# Internationalization
# ---------------------------------------------------------------------------

i18n = gr.I18n(
    en={
        "hero": (
            "# ✨ Fluent Emoji Generator\n"
            "Create colorful, rounded 3D emoji-style icons with an SDXL LoRA.\n\n"
            "**Tip:** English prompts usually produce the most reliable results."
        ),
        "prompt": "Prompt",
        "prompt_placeholder": "Describe one central object…",
        "negative_prompt": "Negative prompt",
        "settings": "Generation settings",
        "lora_weight": "LoRA weight",
        "steps": "Steps",
        "guidance": "Guidance scale",
        "seed": "Seed",
        "seed_info": "Use -1 for a random seed.",
        "generate": "Generate",
        "clear": "Clear",
        "examples": "### Prompt examples",
        "generated_image": "Generated PNG",
        "transparent_image": "Transparent PNG",
        "seed_used": "Seed used",
        "effective_prompt": "Prompt used",
        "download_original": "Download PNG",
        "download_transparent": "Download transparent PNG",
        "about": (
            "The transparent version removes a white or nearly white background. "
            "It works best when the prompt includes `isolated on white background`."
        ),
        "disclaimer": (
            "Independent community project. Not affiliated with or endorsed by Microsoft."
        ),
    },
)


# ---------------------------------------------------------------------------
# Load the pipeline at module level.
# ZeroGPU optimizes CUDA placement performed during Space startup.
# ---------------------------------------------------------------------------

torch.backends.cuda.matmul.allow_tf32 = True

pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)

pipe.load_lora_weights(
    LORA_MODEL_ID,
    weight_name=LORA_WEIGHT_NAME,
    adapter_name=LORA_ADAPTER_NAME,
)

pipe.set_adapters(LORA_ADAPTER_NAME, adapter_weights=0.9)
pipe.set_progress_bar_config(disable=True)
pipe.enable_vae_tiling()
pipe.to("cuda")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def normalize_prompt(prompt: str) -> str:
    """Validate the prompt and add the trigger phrase when it is missing."""
    cleaned = " ".join((prompt or "").strip().split())

    if not cleaned:
        raise gr.Error("Enter a prompt.")

    if len(cleaned) > MAX_PROMPT_LENGTH:
        raise gr.Error(
            f"The prompt is too long (maximum {MAX_PROMPT_LENGTH} characters)."
        )

    if TRIGGER_PHRASE.lower() not in cleaned.lower():
        cleaned = f"{TRIGGER_PHRASE}, {cleaned}"

    return cleaned


def remove_white_background(
    image: Image.Image,
    fully_transparent_from: int = 250,
    feather_from: int = 220,
) -> Image.Image:
    """
    Convert a white or nearly white background to transparency.

    Pixels with all RGB channels at or above fully_transparent_from become
    transparent. Pixels between feather_from and fully_transparent_from receive
    a soft alpha transition to reduce visible white edges.
    """
    rgba = np.asarray(image.convert("RGBA")).copy()
    rgb = rgba[..., :3].astype(np.int16)

    # The darkest RGB channel is used so that saturated light colors remain
    # more opaque than neutral white pixels.
    min_channel = rgb.min(axis=2)

    alpha = np.full(min_channel.shape, 255, dtype=np.uint8)

    transparent_mask = min_channel >= fully_transparent_from
    feather_mask = (
        (min_channel >= feather_from)
        & (min_channel < fully_transparent_from)
    )

    alpha[transparent_mask] = 0

    feather_values = (
        (fully_transparent_from - min_channel[feather_mask])
        / (fully_transparent_from - feather_from)
        * 255
    )
    alpha[feather_mask] = np.clip(feather_values, 0, 255).astype(np.uint8)

    # Respect any alpha channel already present in the source image.
    original_alpha = rgba[..., 3].astype(np.uint16)
    rgba[..., 3] = (
        alpha.astype(np.uint16) * original_alpha // 255
    ).astype(np.uint8)

    return Image.fromarray(rgba, mode="RGBA")


MAX_OUTPUT_FILE_AGE_SECONDS = 3600  # 1 hour


def _cleanup_old_outputs(output_dir: Path) -> None:
    """Delete previously generated files older than MAX_OUTPUT_FILE_AGE_SECONDS.

    Generated PNGs are never referenced again once served to the client, but
    were previously left on disk forever, which would eventually fill up the
    machine's temp storage under sustained use. This keeps only recent files.
    """
    cutoff = time.time() - MAX_OUTPUT_FILE_AGE_SECONDS
    for existing_file in output_dir.glob("fluent_emoji_*.png"):
        try:
            if existing_file.stat().st_mtime < cutoff:
                existing_file.unlink(missing_ok=True)
        except OSError:
            # Best-effort cleanup; never let this break image generation.
            continue


def save_outputs(
    image: Image.Image,
    transparent_image: Image.Image,
    seed: int,
) -> tuple[str, str]:
    """Save both PNG files and return their temporary paths."""
    output_dir = Path(tempfile.gettempdir()) / "fluent_emoji_generator"
    output_dir.mkdir(parents=True, exist_ok=True)
    _cleanup_old_outputs(output_dir)

    unique_id = uuid.uuid4().hex[:10]
    original_path = output_dir / f"fluent_emoji_{seed}_{unique_id}.png"
    transparent_path = (
        output_dir / f"fluent_emoji_{seed}_{unique_id}_transparent.png"
    )

    image.save(original_path, format="PNG")
    transparent_image.save(transparent_path, format="PNG")

    return str(original_path), str(transparent_path)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
@torch.inference_mode()
def generate(
    prompt: str,
    negative_prompt: str,
    lora_weight: float,
    steps: int,
    guidance_scale: float,
    seed: int | float,
):
    effective_prompt = normalize_prompt(prompt)

    negative_prompt = (negative_prompt or "").strip()
    if len(negative_prompt) > MAX_NEGATIVE_PROMPT_LENGTH:
        raise gr.Error(
            f"The negative prompt is too long (maximum "
            f"{MAX_NEGATIVE_PROMPT_LENGTH} characters)."
        )

    # Every numeric input below is re-validated here because the `generate`
    # API endpoint can be called directly, bypassing the slider limits in
    # the UI. Anything malformed raises gr.Error instead of crashing or
    # silently running with an out-of-range / expensive value.
    try:
        seed = int(seed)
    except (TypeError, ValueError):
        raise gr.Error("Seed must be an integer (use -1 for random).")

    if seed < 0:
        actual_seed = random.SystemRandom().randint(MIN_SEED, MAX_SEED)
    elif MIN_SEED <= seed <= MAX_SEED:
        actual_seed = seed
    else:
        raise gr.Error(f"Seed must be between {MIN_SEED} and {MAX_SEED}, or -1.")

    try:
        steps = int(steps)
    except (TypeError, ValueError):
        raise gr.Error("Steps must be an integer.")
    if not (MIN_STEPS <= steps <= MAX_STEPS):
        raise gr.Error(f"Steps must be between {MIN_STEPS} and {MAX_STEPS}.")

    try:
        lora_weight = float(lora_weight)
    except (TypeError, ValueError):
        raise gr.Error("LoRA weight must be a number.")
    if not (MIN_LORA_WEIGHT <= lora_weight <= MAX_LORA_WEIGHT):
        raise gr.Error(
            f"LoRA weight must be between {MIN_LORA_WEIGHT} and {MAX_LORA_WEIGHT}."
        )

    try:
        guidance_scale = float(guidance_scale)
    except (TypeError, ValueError):
        raise gr.Error("Guidance scale must be a number.")
    if not (MIN_GUIDANCE <= guidance_scale <= MAX_GUIDANCE):
        raise gr.Error(
            f"Guidance scale must be between {MIN_GUIDANCE} and {MAX_GUIDANCE}."
        )

    pipe.set_adapters(
        LORA_ADAPTER_NAME,
        adapter_weights=lora_weight,
    )

    generator = torch.Generator(device="cuda").manual_seed(actual_seed)

    result = pipe(
        prompt=effective_prompt,
        negative_prompt=negative_prompt or None,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        width=1024,
        height=1024,
        generator=generator,
        num_images_per_prompt=1,
    )

    image = result.images[0].convert("RGB")
    transparent_image = remove_white_background(image)

    original_path, transparent_path = save_outputs(
        image=image,
        transparent_image=transparent_image,
        seed=actual_seed,
    )

    return (
        image,
        transparent_image,
        actual_seed,
        effective_prompt,
        original_path,
        transparent_path,
    )


def reset_form():
    return (
        DEFAULT_PROMPT,
        DEFAULT_NEGATIVE_PROMPT,
        0.9,
        30,
        7.0,
        -1,
        None,
        None,
        None,
        "",
        None,
        None,
    )


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
.gradio-container {
    max-width: 1180px !important;
    margin: 0 auto !important;
    padding-top: 24px !important;
}
.hero {
    text-align: center;
    margin-bottom: 12px;
}
.hero h1 {
    font-size: clamp(2rem, 5vw, 3.3rem);
    margin-bottom: 0.45rem;
}
.generate-btn {
    min-height: 48px;
    font-weight: 700;
}
.output-image img {
    object-fit: contain !important;
}
.small-note {
    opacity: 0.78;
    font-size: 0.92rem;
}
footer {
    opacity: 0.8;
}
"""

theme = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="violet",
    neutral_hue="slate",
    radius_size="lg",
)


with gr.Blocks(
    theme=theme,
    css=CSS,
    title="Fluent Emoji Generator",
) as demo:
    gr.Markdown(i18n("hero"), elem_classes=["hero"])

    with gr.Row(equal_height=False):
        with gr.Column(scale=5, min_width=320):
            prompt = gr.Textbox(
                value=DEFAULT_PROMPT,
                label=i18n("prompt"),
                placeholder=i18n("prompt_placeholder"),
                lines=4,
                max_lines=8,
                max_length=700,
                autofocus=True,
                buttons=["copy"],
            )

            negative_prompt = gr.Textbox(
                value=DEFAULT_NEGATIVE_PROMPT,
                label=i18n("negative_prompt"),
                lines=2,
                max_lines=5,
                max_length=500,
            )

            with gr.Accordion(i18n("settings"), open=True):
                lora_weight = gr.Slider(
                    minimum=0.0,
                    maximum=1.5,
                    value=0.9,
                    step=0.05,
                    label=i18n("lora_weight"),
                )

                steps = gr.Slider(
                    minimum=15,
                    maximum=50,
                    value=30,
                    step=1,
                    label=i18n("steps"),
                )

                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=12.0,
                    value=7.0,
                    step=0.5,
                    label=i18n("guidance"),
                )

                seed = gr.Number(
                    value=-1,
                    precision=0,
                    label=i18n("seed"),
                    info=i18n("seed_info"),
                )

            with gr.Row():
                generate_button = gr.Button(
                    i18n("generate"),
                    variant="primary",
                    elem_classes=["generate-btn"],
                    scale=3,
                )
                clear_button = gr.Button(
                    i18n("clear"),
                    variant="secondary",
                    scale=1,
                )

            gr.Markdown(i18n("examples"))
            gr.Examples(
                examples=EXAMPLES,
                inputs=[prompt],
                cache_examples=False,
                examples_per_page=6,
            )

        with gr.Column(scale=6, min_width=360):
            with gr.Tabs():
                with gr.Tab(i18n("generated_image")):
                    generated_image = gr.Image(
                        label=i18n("generated_image"),
                        type="pil",
                        format="png",
                        interactive=False,
                        height=540,
                        elem_classes=["output-image"],
                    )

                with gr.Tab(i18n("transparent_image")):
                    transparent_image = gr.Image(
                        label=i18n("transparent_image"),
                        type="pil",
                        format="png",
                        interactive=False,
                        height=540,
                        elem_classes=["output-image"],
                    )

            with gr.Row():
                seed_used = gr.Number(
                    label=i18n("seed_used"),
                    precision=0,
                    interactive=False,
                )

            effective_prompt = gr.Textbox(
                label=i18n("effective_prompt"),
                lines=2,
                interactive=False,
                buttons=["copy"],
            )

            with gr.Row():
                original_file = gr.File(
                    label=i18n("download_original"),
                    interactive=False,
                )
                transparent_file = gr.File(
                    label=i18n("download_transparent"),
                    interactive=False,
                )

            gr.Markdown(i18n("about"), elem_classes=["small-note"])

    gr.Markdown("---")
    gr.Markdown(
        i18n("disclaimer"),
        elem_classes=["small-note"],
    )

    inference_inputs = [
        prompt,
        negative_prompt,
        lora_weight,
        steps,
        guidance_scale,
        seed,
    ]

    inference_outputs = [
        generated_image,
        transparent_image,
        seed_used,
        effective_prompt,
        original_file,
        transparent_file,
    ]

    generate_button.click(
        fn=generate,
        inputs=inference_inputs,
        outputs=inference_outputs,
        api_name="generate",
        concurrency_limit=1,
        concurrency_id="sdxl_gpu",
    )

    prompt.submit(
        fn=generate,
        inputs=inference_inputs,
        outputs=inference_outputs,
        concurrency_limit=1,
        concurrency_id="sdxl_gpu",
    )

    clear_button.click(
        fn=reset_form,
        inputs=None,
        outputs=[
            prompt,
            negative_prompt,
            lora_weight,
            steps,
            guidance_scale,
            seed,
            generated_image,
            transparent_image,
            seed_used,
            effective_prompt,
            original_file,
            transparent_file,
        ],
        queue=False,
    )


if __name__ == "__main__":
    demo.queue(max_size=20).launch(i18n=i18n)