app.py

import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
from typing import Iterable
from gradio.themes import Base
from gradio.themes.utils import colors, fonts, sizes

# ── Windows 98 Color Palettes ───────────────────────────────────────────

colors.win98_blue = colors.Color(
    name="win98_blue",
    c50="#E6E6F5",
    c100="#CCCCEB",
    c200="#9999D6",
    c300="#6666C2",
    c400="#3333AD",
    c500="#000080",
    c600="#000073",
    c700="#000066",
    c800="#000059",
    c900="#00004D",
    c950="#000040",
)

colors.win98_gray = colors.Color(
    name="win98_gray",
    c50="#F5F5F5",
    c100="#EBEBEB",
    c200="#D6D6D6",
    c300="#C0C0C0",
    c400="#A8A8A8",
    c500="#808080",
    c600="#696969",
    c700="#555555",
    c800="#404040",
    c900="#2D2D2D",
    c950="#1A1A1A",
)


# ── Windows 98 Theme ────────────────────────────────────────────────────

class Windows98Theme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.win98_gray,
        secondary_hue: colors.Color | str = colors.win98_blue,
        neutral_hue: colors.Color | str = colors.win98_gray,
        text_size: sizes.Size | str = sizes.text_md,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            "Tahoma", "MS Sans Serif", "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            "Courier New", "Fixedsys", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            # ── Body ──
            body_background_fill="#c0c0c0",
            body_background_fill_dark="#c0c0c0",
            background_fill_primary="#c0c0c0",
            background_fill_primary_dark="#c0c0c0",
            background_fill_secondary="#d4d0c8",
            background_fill_secondary_dark="#d4d0c8",

            # ── Blocks ──
            block_background_fill="#c0c0c0",
            block_border_width="2px",
            block_shadow="none",
            block_title_text_weight="bold",
            block_label_background_fill="#c0c0c0",

            # ── Primary Button (Navy Blue) ──
            button_primary_background_fill="#000080",
            button_primary_background_fill_hover="#0000a0",
            button_primary_background_fill_dark="#000080",
            button_primary_background_fill_hover_dark="#0000a0",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_shadow="none",

            # ── Secondary Button (Gray) ──
            button_secondary_background_fill="#c0c0c0",
            button_secondary_background_fill_hover="#d4d0c8",
            button_secondary_background_fill_dark="#c0c0c0",
            button_secondary_background_fill_hover_dark="#d4d0c8",
            button_secondary_text_color="#000000",
            button_secondary_text_color_hover="#000000",

            button_large_padding="6px 16px",

            # ── Slider ──
            slider_color="#000080",
            slider_color_dark="#000080",

            # ── Accent ──
            color_accent_soft="#d4d0c8",

            # ── Input ──
            input_background_fill="#ffffff",
            input_background_fill_dark="#ffffff",

            # ── Text ──
            body_text_color="#000000",
            body_text_color_dark="#000000",
        )


windows_98_theme = Windows98Theme()

# ── Device & Model Setup ────────────────────────────────────────────────

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

print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("Using device:", device)

from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

dtype = torch.bfloat16

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.0",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
        torch_dtype=dtype,
        device_map='cuda'
    ),
    torch_dtype=dtype
).to(device)

try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Flash Attention 3 Processor set successfully.")
except Exception as e:
    print(f"Warning: Could not set FA3 processor: {e}")

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


# ── Helper Functions ────────────────────────────────────────────────────

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024

    original_width, original_height = image.size

    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)

    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8

    return new_width, new_height


@spaces.GPU
def infer(
    images,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True)
):
    gc.collect()
    torch.cuda.empty_cache()

    if not images:
        raise gr.Error("Please upload at least one image to edit.")

    pil_images = []
    if images is not None:
        for item in images:
            try:
                if isinstance(item, tuple) or isinstance(item, list):
                    path_or_img = item[0]
                else:
                    path_or_img = item

                if isinstance(path_or_img, str):
                    pil_images.append(Image.open(path_or_img).convert("RGB"))
                elif isinstance(path_or_img, Image.Image):
                    pil_images.append(path_or_img.convert("RGB"))
                else:
                    pil_images.append(Image.open(path_or_img.name).convert("RGB"))
            except Exception as e:
                print(f"Skipping invalid image item: {e}")
                continue

    if not pil_images:
        raise gr.Error("Could not process uploaded images.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"

    width, height = update_dimensions_on_upload(pil_images[0])

    try:
        result_image = pipe(
            image=pil_images,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]

        return result_image, seed

    except Exception as e:
        raise e
    finally:
        gc.collect()
        torch.cuda.empty_cache()


@spaces.GPU
def infer_example(images, prompt):
    if not images:
        return None, 0

    if isinstance(images, str):
        images_list = [images]
    else:
        images_list = images

    result, seed = infer(
        images=images_list,
        prompt=prompt,
        seed=0,
        randomize_seed=True,
        guidance_scale=1.0,
        steps=4
    )
    return result, seed


# ── CSS ──────────────────────────────────────────────────────────────────

css = """
/* ═══════════════════════════════════════════════
   WINDOWS 98 RETRO THEME
   ═══════════════════════════════════════════════ */

/* ── Global Reset ── */
*, *::before, *::after {
    border-radius: 0px !important;
}

.gradio-container {
    background: #c0c0c0 !important;
}

/* ── Window Frame ── */
#col-container {
    margin: 0 auto;
    max-width: 1000px;
    background: #c0c0c0;
    border-top: 2px solid #ffffff;
    border-left: 2px solid #ffffff;
    border-right: 2px solid #404040;
    border-bottom: 2px solid #404040;
    box-shadow: 1px 1px 0 0 #000000;
    padding: 2px;
}

/* ── Title Bar ── */
#main-title {
    background: linear-gradient(90deg, #000080, #1084d0) !important;
    padding: 3px 8px !important;
    margin: 0 !important;
    border: none !important;
    box-shadow: none !important;
    position: relative;
}

#main-title h1 {
    font-size: 13px !important;
    color: white !important;
    font-weight: bold !important;
    margin: 0 !important;
    padding: 1px 0 !important;
    text-shadow: 1px 1px 0px rgba(0,0,0,0.5);
    font-family: "Tahoma", "MS Sans Serif", sans-serif !important;
}

/* Faux Window Controls */
#main-title::after {
    content: "— □ ✕";
    position: absolute;
    right: 4px;
    top: 3px;
    color: #000000;
    font-size: 9px;
    letter-spacing: 2px;
    background: #c0c0c0;
    border-top: 2px solid #ffffff;
    border-left: 2px solid #ffffff;
    border-right: 2px solid #808080;
    border-bottom: 2px solid #808080;
    padding: 0 5px;
    font-family: "Tahoma", sans-serif;
    line-height: 15px;
}

/* ── Description / Menu Bar ── */
#desc-text {
    background: #c0c0c0 !important;
    border: none !important;
    border-bottom: 1px solid #808080 !important;
    padding: 4px 6px !important;
    margin: 0 0 4px 0 !important;
    box-shadow: none !important;
}

#desc-text p {
    font-size: 11px !important;
    margin: 2px 0 !important;
    color: #000000 !important;
}

#desc-text a {
    color: #000080 !important;
    text-decoration: underline !important;
}

/* ── Raised Panels / Blocks ── */
.block, .form {
    border-top: 2px solid #ffffff !important;
    border-left: 2px solid #ffffff !important;
    border-right: 2px solid #808080 !important;
    border-bottom: 2px solid #808080 !important;
    background: #c0c0c0 !important;
    box-shadow: none !important;
}

/* ── Sunken Text Inputs ── */
textarea,
input[type="text"],
input[type="number"],
input[type="search"] {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    background: #ffffff !important;
    box-shadow: none !important;
    font-family: "Tahoma", "MS Sans Serif", sans-serif !important;
    outline: none !important;
}

textarea:focus,
input[type="text"]:focus,
input[type="number"]:focus {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    box-shadow: none !important;
    outline: none !important;
}

/* ── Sunken Image / Gallery Frames ── */
.image-frame {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    background: #ffffff !important;
}

/* ── Upload Area ── */
.upload-area {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    background: #ffffff !important;
}

/* ── Buttons (Raised) ── */
button {
    border-top: 2px solid #ffffff !important;
    border-left: 2px solid #ffffff !important;
    border-right: 2px solid #808080 !important;
    border-bottom: 2px solid #808080 !important;
    box-shadow: none !important;
    font-family: "Tahoma", "MS Sans Serif", sans-serif !important;
}

button:active {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
}

/* ── Checkbox (Sunken) ── */
input[type="checkbox"] {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    background: #ffffff !important;
}

/* ── Slider ── */
input[type="range"] {
    border: none !important;
    background: transparent !important;
}

/* ── Select / Dropdown (Sunken) ── */
select {
    border-top: 2px solid #808080 !important;
    border-left: 2px solid #808080 !important;
    border-right: 2px solid #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
    background: #ffffff !important;
}

/* ── Accordion (Group Box) ── */
.accordion {
    border: 2px groove #c0c0c0 !important;
    box-shadow: none !important;
}

/* ── Labels ── */
.block-label, .label-wrap {
    background: #c0c0c0 !important;
}

/* ── Tab Nav (if present) ── */
.tab-nav button {
    border-top: 2px solid #ffffff !important;
    border-left: 2px solid #ffffff !important;
    border-right: 2px solid #808080 !important;
    border-bottom: none !important;
    background: #c0c0c0 !important;
    margin-right: 1px !important;
}

.tab-nav button.selected {
    border-bottom: 2px solid #c0c0c0 !important;
    position: relative;
    z-index: 2;
}

/* ── Scrollbars (Classic Win98) ── */
::-webkit-scrollbar {
    width: 16px;
    height: 16px;
}

::-webkit-scrollbar-track {
    background: #c0c0c0;
}

::-webkit-scrollbar-thumb {
    background: #c0c0c0;
    border-top: 2px solid #ffffff;
    border-left: 2px solid #ffffff;
    border-right: 2px solid #808080;
    border-bottom: 2px solid #808080;
}

::-webkit-scrollbar-button {
    display: block;
    height: 16px;
    width: 16px;
    background: #c0c0c0;
    border-top: 2px solid #ffffff;
    border-left: 2px solid #ffffff;
    border-right: 2px solid #808080;
    border-bottom: 2px solid #808080;
}

::-webkit-scrollbar-corner {
    background: #c0c0c0;
}

/* ── Progress Bar ── */
.progress-bar {
    background: #000080 !important;
}

/* ── Toast / Error Dialogs ── */
.toast-wrap, .toast-body {
    border-top: 2px solid #ffffff !important;
    border-left: 2px solid #ffffff !important;
    border-right: 2px solid #808080 !important;
    border-bottom: 2px solid #808080 !important;
    background: #c0c0c0 !important;
    color: #000000 !important;
}

/* ── Footer ── */
footer {
    border-top: 1px solid #808080 !important;
    background: #c0c0c0 !important;
}
"""

# ── Gradio Interface ─────────────────────────────────────────────────────

with gr.Blocks(css=css, theme=windows_98_theme) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **FireRed-Image-Edit-1.0-Fast**", elem_id="main-title")
        gr.Markdown(
            "Perform image edits using "
            "[FireRed-Image-Edit-1.0](https://huggingface.co/FireRedTeam/FireRed-Image-Edit-1.0) "
            "with 4-step fast inference.",
            elem_id="desc-text"
        )

        with gr.Row(equal_height=True):
            with gr.Column():
                images = gr.Gallery(
                    label="Upload Images",
                    type="filepath",
                    columns=2,
                    rows=1,
                    height=300,
                    allow_preview=True
                )

                prompt = gr.Text(
                    label="Edit Prompt",
                    show_label=True,
                    placeholder="e.g., transform into anime, upscale, change lighting...",
                )

                run_button = gr.Button("Edit Image", variant="primary")

            with gr.Column():
                output_image = gr.Image(
                    label="Output Image",
                    interactive=False,
                    format="png",
                    height=395
                )

                with gr.Accordion("Advanced Settings", open=False, visible=False):
                    seed = gr.Slider(
                        label="Seed", minimum=0,
                        maximum=MAX_SEED, step=1, value=0
                    )
                    randomize_seed = gr.Checkbox(
                        label="Randomize Seed", value=True
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance Scale", minimum=1.0,
                        maximum=10.0, step=0.1, value=1.0
                    )
                    steps = gr.Slider(
                        label="Inference Steps", minimum=1,
                        maximum=50, step=1, value=4
                    )

    run_button.click(
        fn=infer,
        inputs=[images, prompt, seed, randomize_seed, guidance_scale, steps],
        outputs=[output_image, seed]
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(
        mcp_server=True,
        ssr_mode=False,
        show_error=True
    )