app.py

import importlib
import random

# Import spaces before torch so ZeroGPU can patch CUDA handling at startup.
spaces = None
try:
    spaces = importlib.import_module("spaces")
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

import gradio as gr  # noqa: E402
import torch  # noqa: E402

# Patch transformers to work around the Qwen3-VL rope_scaling bug
try:
    from transformers.models.qwen3_vl.modeling_qwen3_vl import (
        Qwen3VLTextRotaryEmbedding,
    )

    old_init = Qwen3VLTextRotaryEmbedding.__init__

    def patched_init(self, config, *args, **kwargs):
        if getattr(config, "rope_scaling", None) is None:
            config.rope_scaling = {
                "mrope_section": [24, 20, 20],
                "rope_type": "default",
            }
        old_init(self, config, *args, **kwargs)

    Qwen3VLTextRotaryEmbedding.__init__ = patched_init
    print("Successfully patched Qwen3VLTextRotaryEmbedding.__init__")
except Exception as patch_err:
    print(f"Skipping transformers patch: {patch_err}")

# Patch gradio_client to work around the json_schema_to_python_type boolean schema bug
try:
    import gradio_client.utils as gr_utils

    old_json_schema = gr_utils._json_schema_to_python_type

    def patched_json_schema(schema, *args, **kwargs):
        if isinstance(schema, bool):
            return "Any"
        return old_json_schema(schema, *args, **kwargs)

    gr_utils._json_schema_to_python_type = patched_json_schema
    print("Successfully patched gradio_client._json_schema_to_python_type")
except Exception as patch_err:
    print(f"Skipping gradio_client patch: {patch_err}")

from diffusers import Krea2Pipeline  # noqa: E402

# Select device and torch dtype. On ZeroGPU, importing spaces before torch makes
# CUDA appear available while keeping the model packable at module scope.
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

print(f"Loading Krea-2-Turbo model. Device: {device}, Dtype: {dtype}")

# Load the tokenizer explicitly as Qwen2TokenizerFast to bypass slow tokenizer missing vocab files and list AttributeError
from transformers import Qwen2TokenizerFast  # noqa: E402

tokenizer = Qwen2TokenizerFast.from_pretrained(
    "krea/Krea-2-Turbo", subfolder="tokenizer", extra_special_tokens={}
)

# Load the model at the module level passing the pre-loaded fast tokenizer
# In HF ZeroGPU, CUDA emulation is active on startup, allowing .to("cuda") to succeed
pipe = Krea2Pipeline.from_pretrained(
    "krea/Krea-2-Turbo", tokenizer=tokenizer, torch_dtype=dtype, use_safetensors=True
).to(device)

# Custom CSS for a light-only polished UI. The dark-mode selectors are
# intentional: they override Gradio/browser dark mode when the system is dark.
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700;800&display=swap');

:root,
html,
body,
.gradio-container,
.gradio-container.dark {
    color-scheme: light !important;
    --body-background-fill: #f8fafc !important;
    --body-text-color: #0f172a !important;
    --block-background-fill: rgba(255, 255, 255, 0.86) !important;
    --block-border-color: rgba(148, 163, 184, 0.28) !important;
    --input-background-fill: #ffffff !important;
    --input-border-color: #cbd5e1 !important;
    --input-placeholder-color: #94a3b8 !important;
    --neutral-50: #f8fafc !important;
    --neutral-100: #f1f5f9 !important;
    --neutral-200: #e2e8f0 !important;
    --neutral-700: #334155 !important;
    --neutral-800: #1e293b !important;
    --neutral-900: #0f172a !important;
}

body, .gradio-container, .gradio-container.dark {
    font-family: 'Outfit', sans-serif !important;
    background:
        radial-gradient(circle at top left, rgba(216, 180, 254, 0.46), transparent 32%),
        radial-gradient(circle at top right, rgba(147, 197, 253, 0.38), transparent 34%),
        linear-gradient(135deg, #ffffff 0%, #f8fafc 45%, #eef2ff 100%) !important;
    color: #0f172a !important;
}

/* App container centering */
.app-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}

/* Main header card */
.header {
    text-align: center;
    margin-bottom: 2.5rem;
    padding: 2.5rem;
    background: rgba(255, 255, 255, 0.78);
    backdrop-filter: blur(18px);
    border: 1px solid rgba(148, 163, 184, 0.28);
    border-radius: 24px;
    box-shadow: 0 18px 50px rgba(99, 102, 241, 0.14);
}

.header h1 {
    font-size: 3.2rem !important;
    font-weight: 800 !important;
    background: linear-gradient(90deg, #c026d3, #7c3aed, #2563eb);
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    margin-bottom: 0.5rem !important;
    letter-spacing: -0.02em;
    text-shadow: 0 10px 30px rgba(124, 58, 237, 0.16);
}

.header p {
    font-size: 1.1rem !important;
    color: #475569 !important;
    max-width: 600px;
    margin: 0 auto !important;
    line-height: 1.6;
}

/* Custom panels */
.glass-panel {
    background: rgba(255, 255, 255, 0.84) !important;
    backdrop-filter: blur(16px) !important;
    border: 1px solid rgba(148, 163, 184, 0.26) !important;
    border-radius: 20px !important;
    box-shadow: 0 14px 38px rgba(15, 23, 42, 0.08) !important;
    padding: 1.5rem !important;
    transition: border-color 0.3s ease, box-shadow 0.3s ease !important;
}

.glass-panel:hover {
    border-color: rgba(124, 58, 237, 0.34) !important;
    box-shadow: 0 18px 45px rgba(124, 58, 237, 0.14) !important;
}

/* Call-to-action generate button */
.btn-generate {
    background: linear-gradient(135deg, #d946ef 0%, #8b5cf6 50%, #6366f1 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    font-weight: 700 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
    padding: 12px 24px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 8px 22px rgba(139, 92, 246, 0.32) !important;
}

.btn-generate:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 12px 28px rgba(139, 92, 246, 0.44) !important;
}

.btn-generate:active {
    transform: translateY(1px) !important;
}

/* Styled inputs and select elements */
.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container.dark input,
.gradio-container.dark textarea,
.gradio-container.dark select {
    background-color: #ffffff !important;
    border: 1px solid #cbd5e1 !important;
    color: #0f172a !important;
    border-radius: 12px !important;
    padding: 10px 14px !important;
    font-size: 0.95rem !important;
    transition: all 0.2s ease !important;
}

.gradio-container input:focus,
.gradio-container textarea:focus,
.gradio-container select:focus {
    border-color: #8b5cf6 !important;
    box-shadow: 0 0 0 2px rgba(139, 92, 246, 0.22) !important;
}

.gradio-container label,
.gradio-container .prose,
.gradio-container .markdown,
.gradio-container.dark label,
.gradio-container.dark .prose,
.gradio-container.dark .markdown {
    color: #1e293b !important;
}

/* Custom License badge styling */
.badge-license {
    display: inline-block;
    padding: 5px 12px;
    border-radius: 100px;
    background: rgba(139, 92, 246, 0.12);
    color: #7c3aed;
    font-size: 0.825rem;
    font-weight: 600;
    border: 1px solid rgba(139, 92, 246, 0.22);
    margin-top: 1rem;
}
"""


def infer(prompt, style, aspect_ratio, steps, guidance_scale, seed):
    # Apply style preset template
    style_prompts = {
        "None": "{prompt}",
        "Cinematic": "{prompt}, cinematic shot, 35mm lens, depth of field, anamorphic flare, professional lighting, photorealistic, 8k",
        "Photorealistic": "{prompt}, photorealistic, ultra-detailed, 8k resolution, highly detailed, dramatic lighting, sharp focus",
        "Digital Art": "{prompt}, digital painting, concept art, sharp focus, vibrant colors, detailed, artstation trending",
        "Fantasy / Mythic": "{prompt}, high fantasy, magical, ethereal, concept art, mystical lighting, detailed, epic composition",
        "Cyberpunk": "{prompt}, cyberpunk aesthetic, neon lights, futuristic city, highly detailed, dark synthwave style, raytracing",
        "Anime / Manga": "{prompt}, anime style, colorful, studio ghibli, detailed background, vibrant, aesthetic key visual, 4k",
        "Vintage Film": "{prompt}, vintage photograph, film grain, muted colors, nostalgic, warm lighting, old photo style",
    }

    formatted_prompt = style_prompts.get(style, "{prompt}").format(prompt=prompt)

    # Handle aspect ratios and resolutions
    aspect_ratios = {
        "1:1 Square (1024x1024)": (1024, 1024),
        "9:16 Vertical (768x1360)": (768, 1360),
        "16:9 Landscape (1360x768)": (1360, 768),
        "4:3 Portrait (896x1152)": (896, 1152),
        "3:4 Landscape (1152x896)": (1152, 896),
    }
    width, height = aspect_ratios.get(aspect_ratio, (1024, 1024))

    # Handle seed selection
    if seed == -1 or seed is None:
        seed = random.randint(0, 2**32 - 1)
    else:
        seed = int(seed)

    generator = torch.Generator(device=device).manual_seed(seed)

    # Run the model
    # Note: Krea-2-Turbo is optimized for 8 inference steps and 0.0 guidance scale.
    image = pipe(
        prompt=formatted_prompt,
        width=width,
        height=height,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        generator=generator,
    ).images[0]

    return image, seed


# Apply ZeroGPU decorator dynamically if spaces library is present
if HAS_SPACES and spaces is not None:
    generate_fn = getattr(spaces, "GPU")(duration=35)(infer)
else:
    generate_fn = infer

# Define Gradio blocks application
theme = gr.themes.Default(primary_hue="purple", secondary_hue="indigo")

with gr.Blocks() as demo:
    with gr.Column(elem_classes="app-container"):
        gr.HTML("""
        <div class="header">
            <h1>⚡ Kreatin</h1>
            <p>Create polished images fast with a clean prompt studio powered by Krea 2 Turbo, Krea's distilled 12B parameter Diffusion Transformer optimized for 8-step generation.</p>
            <span class="badge-license">Powered by Krea 2 Turbo • Krea 2 Community License</span>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=5):
                with gr.Group(elem_classes="glass-panel"):
                    prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Describe the image you want to create...",
                        lines=3,
                        elem_id="prompt-input",
                    )

                    with gr.Row():
                        style = gr.Dropdown(
                            label="Style Preset",
                            choices=[
                                "None",
                                "Cinematic",
                                "Photorealistic",
                                "Digital Art",
                                "Fantasy / Mythic",
                                "Cyberpunk",
                                "Anime / Manga",
                                "Vintage Film",
                            ],
                            value="None",
                        )

                        aspect_ratio = gr.Dropdown(
                            label="Aspect Ratio",
                            choices=[
                                "1:1 Square (1024x1024)",
                                "9:16 Vertical (768x1360)",
                                "16:9 Landscape (1360x768)",
                                "4:3 Portrait (896x1152)",
                                "3:4 Landscape (1152x896)",
                            ],
                            value="1:1 Square (1024x1024)",
                        )

                    generate_btn = gr.Button(
                        "Generate Image", elem_classes="btn-generate"
                    )

                with gr.Accordion(
                    "Advanced Settings", open=False, elem_classes="glass-panel"
                ):
                    with gr.Row():
                        steps = gr.Slider(
                            label="Inference Steps",
                            minimum=1,
                            maximum=25,
                            step=1,
                            value=8,
                            info="Turbo is optimized for 8 steps.",
                        )

                        guidance_scale = gr.Slider(
                            label="Guidance Scale (CFG)",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=0.0,
                            info="Turbo model works best with 0.0 guidance scale.",
                        )

                    seed = gr.Number(
                        label="Seed (-1 for random)", value=-1, precision=0
                    )

            with gr.Column(scale=5):
                with gr.Group(elem_classes="glass-panel"):
                    output_image = gr.Image(
                        label="Generated Output",
                        type="pil",
                        interactive=False,
                    )

                    with gr.Row():
                        used_seed = gr.Number(
                            label="Used Seed", interactive=False, precision=0
                        )

        # Connect button click
        generate_btn.click(
            fn=generate_fn,
            inputs=[prompt, style, aspect_ratio, steps, guidance_scale, seed],
            outputs=[output_image, used_seed],
        )

        # Connect prompt enter key
        prompt.submit(
            fn=generate_fn,
            inputs=[prompt, style, aspect_ratio, steps, guidance_scale, seed],
            outputs=[output_image, used_seed],
        )

        gr.HTML("<br>")

        # Examples section
        gr.Examples(
            examples=[
                [
                    "A stunning close-up portrait of a majestic lion with glowing cyan eyes, digital art, highly detailed",
                    "Digital Art",
                    "1:1 Square (1024x1024)",
                    8,
                    0.0,
                    42,
                ],
                [
                    "A futuristic cyberpunk street at night, neon reflections in the puddles, rain, cinematic lighting, 8k",
                    "Cyberpunk",
                    "16:9 Landscape (1360x768)",
                    8,
                    0.0,
                    1337,
                ],
                [
                    "An enchanted cottage in the middle of a mystical forest, fireflies, high fantasy, warm volumetric light",
                    "Fantasy / Mythic",
                    "4:3 Portrait (896x1152)",
                    8,
                    0.0,
                    999,
                ],
                [
                    "A high-fashion editorial photo of a model in a vibrant flowing dress, architectural backdrop, cinematic lighting",
                    "Photorealistic",
                    "9:16 Vertical (768x1360)",
                    8,
                    0.0,
                    777,
                ],
                [
                    "Retro 90s anime style screenshot of a pilot in a spaceship cockpit looking at the stars",
                    "Anime / Manga",
                    "16:9 Landscape (1360x768)",
                    8,
                    0.0,
                    12345,
                ],
            ],
            inputs=[prompt, style, aspect_ratio, steps, guidance_scale, seed],
            outputs=[output_image, used_seed],
            fn=generate_fn,
            cache_examples=False,
        )

if __name__ == "__main__":
    demo.queue().launch(css=custom_css, theme=theme)