app.py

#!/usr/bin/env python
"""
GritAI — AI Image Generator  (Project #6 of "30 AI Projects in 15 Days")

FREE public tier: text -> image on Hugging Face ZeroGPU using FLUX.1-schnell
(Apache-2.0, guidance-distilled, 4-step). Fast, $0/image, always-on — the top of the funnel.

PRO tier ("GritAI Studio"): the same prompt routed to our OWN full stack (local ComfyUI on an
RTX 5070 Ti) for custom checkpoints/LoRAs, upscaling to 4K, batch, and style packs. Shown here as
a waitlist — we light up the local backend once there's demand.

Runs on a Gradio ZeroGPU Space. The heavy work sits inside a @spaces.GPU function so the GPU is
only held for the seconds a generation actually takes.
"""
import os
import random

import gradio as gr

# --- ZeroGPU decorator, with a no-op shim so the file imports/builds without HF infra (local smoke test) ---
try:
    import spaces  # provided by the ZeroGPU runtime
    _HAS_SPACES = True
except Exception:  # pragma: no cover - local dev without the spaces package
    _HAS_SPACES = False

    class spaces:  # type: ignore
        @staticmethod
        def GPU(*a, **k):
            def wrap(fn):
                return fn
            return wrap


MODEL_ID = "black-forest-labs/FLUX.1-schnell"
MAX_SEED = 2**31 - 1

# Aspect presets (Flux likes multiples of 16; keep the megapixel count ~1MP for speed on ZeroGPU)
RATIOS = {
    "Square 1:1 (1024×1024)": (1024, 1024),
    "Portrait 3:4 (896×1152)": (896, 1152),
    "Portrait 9:16 (768×1344)": (768, 1344),
    "Landscape 4:3 (1152×896)": (1152, 896),
    "Landscape 16:9 (1344×768)": (1344, 768),
}

EXAMPLES = [
    "a cozy bookstore cafe at golden hour, warm rim light, shallow depth of field, photoreal",
    "product shot of a matte-black ceramic coffee mug on slate, studio softbox lighting, crisp",
    "a friendly cartoon fox mascot for a coffee brand, flat vector, bold clean lines",
    "isometric 3D miniature of a tiny mountain town, soft clay render, pastel palette",
    "vintage travel poster of Lawton, Oklahoma at sunset, screen-print texture, bold type space",
]

# Pipeline handling. The weights (~24GB) are loaded to CPU at STARTUP on the Space (see the preload
# below) so the download never eats into the per-request ZeroGPU time budget. Locally we never load
# it — the module imports (and the UI builds) without CUDA / diffusers present.
_PIPE = None


def _get_pipe():
    """Return the FLUX pipeline on CPU, loading (downloading) it on first use."""
    global _PIPE
    if _PIPE is None:
        import torch
        from diffusers import FluxPipeline

        _PIPE = FluxPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
    return _PIPE


# Preload weights at startup, but ONLY on the real ZeroGPU runtime — never during local UI preview.
if _HAS_SPACES:
    try:
        _get_pipe()
        print("[startup] FLUX.1-schnell weights loaded to CPU; GPU attaches per request", flush=True)
    except Exception as _e:  # keep the Space up even if preload hiccups; first call will retry
        print(f"[startup] pipe preload failed (will retry on first request): {_e}", flush=True)


@spaces.GPU(duration=75)
def generate(prompt, ratio, steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt to generate an image.")
    import torch

    # On the ZeroGPU Space a GPU is attached while this runs. If there's no CUDA device, we're in a
    # local UI preview — fail with a friendly message instead of a model-load traceback.
    if not torch.cuda.is_available():
        raise gr.Error(
            "🖼️ Local UI preview — image generation runs on the ZeroGPU Space (needs a GPU). "
            "The controls, examples, and Studio waitlist here are all live."
        )

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    width, height = RATIOS.get(ratio, (1024, 1024))

    try:
        pipe = _get_pipe().to("cuda")
        generator = torch.Generator(device="cuda").manual_seed(seed)
        image = pipe(
            prompt=prompt.strip(),
            width=width,
            height=height,
            num_inference_steps=int(steps),
            guidance_scale=0.0,      # schnell is guidance-distilled — CFG/negative prompt don't apply
            max_sequence_length=256,
            generator=generator,
        ).images[0]
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Generation failed — please try again. ({type(e).__name__})")
    return image, seed


# ----------------------------------------------------------------------------- UI
THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.orange, neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
).set(
    body_background_fill="#0a0b0e", body_text_color="#eef0f4", body_text_color_subdued="#9aa1b0",
    background_fill_primary="#111318", background_fill_secondary="#0e1015",
    block_background_fill="#111318", block_border_color="#20242e",
    block_label_background_fill="#0e1015", block_label_text_color="#9aa1b0", block_title_text_color="#eef0f4",
    border_color_primary="#20242e", border_color_accent="#c9451f",
    input_background_fill="#0e1015", input_border_color="#2b303c", input_border_color_focus="#c9451f",
    button_primary_background_fill="#ff5a2a", button_primary_background_fill_hover="#ff7d54",
    button_primary_text_color="#140a06", button_primary_border_color="#c9451f",
    button_secondary_background_fill="#15181f", button_secondary_text_color="#eef0f4",
    button_secondary_border_color="#2b303c",
    color_accent="#ff5a2a", color_accent_soft="#2b1810", slider_color="#ff5a2a",
    checkbox_background_color_selected="#ff5a2a", checkbox_border_color_focus="#c9451f",
)

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Archivo:wght@700;800;900&family=IBM+Plex+Mono:wght@400;500;600&family=IBM+Plex+Sans:wght@400;500;600&display=swap');
.gradio-container{background:#0a0b0e !important}
.gradio-container::before{content:"";position:fixed;inset:0;z-index:0;pointer-events:none;
 background:repeating-linear-gradient(90deg,transparent 0 63px,rgba(255,255,255,.018) 63px 64px),
  repeating-linear-gradient(0deg,transparent 0 63px,rgba(255,255,255,.018) 63px 64px),
  radial-gradient(120% 70% at 50% -10%,#141821 0%,#0a0b0e 55%)}
.main,.contain{position:relative;z-index:1}
#ghead{display:flex;align-items:center;gap:14px;padding:6px 2px 4px}
#ghead .co{font-family:'Archivo',sans-serif;font-weight:900;font-size:19px;letter-spacing:.5px;color:#eef0f4;line-height:1}
#ghead .co b{color:#ff5a2a}
#ghead .sub{font-family:'IBM Plex Mono',monospace;font-size:10.5px;letter-spacing:.28em;text-transform:uppercase;color:#5f6675;margin-top:5px}
#ghead .tags{margin-left:auto;display:flex;gap:8px}
#ghead .tag{font-family:'IBM Plex Mono',monospace;font-size:10px;letter-spacing:.14em;text-transform:uppercase;color:#9aa1b0;border:1px solid #2b303c;border-radius:2px;padding:5px 9px}
#ghead .tag.hot{color:#ff5a2a;border-color:#c9451f}
#ghero h1{font-family:'Archivo',sans-serif;font-weight:800;font-size:clamp(26px,4.2vw,40px);line-height:1.05;letter-spacing:-.02em;margin:14px 0 4px;color:#eef0f4}
#ghero h1 .g{color:#ff5a2a}
#ghero p{color:#9aa1b0;margin:0;font-size:14.5px}
footer{visibility:hidden}
button.selected{color:#ff5a2a !important}
.gritai-foot{text-align:center;color:#5f6675;font-size:.8rem;margin-top:16px;font-family:'IBM Plex Mono',monospace;letter-spacing:.05em}
.gritai-foot a{color:#9aa1b0;text-decoration:none}
.pro-card{border:1px solid #2b303c;border-left:2px solid #ff5a2a;border-radius:4px;padding:16px 18px;background:#111318}
.pro-card h3{color:#eef0f4}
"""

with gr.Blocks(theme=THEME, css=CSS, title="GritAI · AI Image Generator") as demo:
    gr.HTML('''<div id="ghead">
      <div style="width:34px;height:34px;flex:none"><svg width="34" height="34" viewBox="0 0 38 38" fill="none"><rect x="1" y="1" width="36" height="36" rx="2" stroke="#2b303c"/><path d="M27 12.5A9 9 0 1 0 28 22H19" stroke="#ff5a2a" stroke-width="2.4" stroke-linecap="square"/><rect x="10.5" y="10.5" width="3" height="3" fill="#ff5a2a"/></svg></div>
      <div><div class="co">GRIT<b>AI</b></div><div class="sub">AI Image Generator</div></div>
      <div class="tags"><span class="tag hot">&#9656; Live</span><span class="tag">SDVOSB</span><span class="tag">Lawton OK</span></div>
    </div>
    <div id="ghero"><h1>Type an idea. Get an <span class="g">image</span> in seconds.</h1>
    <p>FLUX.1-schnell on free GPU &middot; $0/image &middot; Project #6 of 30 in 15 days.</p></div>''')

    with gr.Tab("Generate (free)"):
        with gr.Row():
            with gr.Column(scale=3):
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="a cozy bookstore cafe at golden hour, photoreal…",
                    lines=3,
                    autofocus=True,
                )
                with gr.Row():
                    ratio = gr.Dropdown(
                        choices=list(RATIOS.keys()),
                        value="Square 1:1 (1024×1024)",
                        label="Aspect ratio",
                    )
                    steps = gr.Slider(1, 8, value=4, step=1, label="Steps (schnell needs only ~4)")
                with gr.Row():
                    seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                    randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
                go = gr.Button("Generate ▶", variant="primary")
                gr.Examples(examples=EXAMPLES, inputs=prompt, label="Try one")
            with gr.Column(scale=4):
                out_img = gr.Image(label="Result", type="pil", height=520, show_download_button=True)
                used_seed = gr.Number(label="Seed used", interactive=False)

        go.click(
            generate,
            inputs=[prompt, ratio, steps, seed, randomize_seed],
            outputs=[out_img, used_seed],
            api_name="generate",
        )
        prompt.submit(
            generate,
            inputs=[prompt, ratio, steps, seed, randomize_seed],
            outputs=[out_img, used_seed],
            api_name=False,
        )

    with gr.Tab("GritAI Studio (Pro)"):
        gr.Markdown(
            "<div class='pro-card'>\n\n"
            "### 🚀 Want more than the free demo?\n"
            "The free tier runs a fast open model on shared GPU. **GritAI Studio** routes your prompt to "
            "our **own production stack** for work you can actually ship:\n\n"
            "- **Custom models & LoRAs** — your brand, product, or character, trained in.\n"
            "- **4K upscaling & detail passes** — print- and web-ready output.\n"
            "- **Batch & style packs** — dozens of on-brand variations in one go.\n"
            "- **Commercial-ready** — consistent characters, product shots, ad creative.\n\n"
            "*Built on the same GPU rig behind GritAI's video & 3D-print pipelines.*\n\n"
            "</div>"
        )
        with gr.Row():
            email = gr.Textbox(label="Email", placeholder="you@business.com", scale=3)
            use_case = gr.Textbox(label="What would you make?", placeholder="e.g. product photos for my shop", scale=4)
        join = gr.Button("Join the Studio waitlist", variant="primary")
        waitlist_msg = gr.Markdown(visible=False)

        def _join(email_v, use_case_v):
            if not email_v or "@" not in email_v:
                return gr.update(value="⚠️ Please enter a valid email.", visible=True)
            # Waitlist capture is wired at deploy time (Scribe/vault). Stub keeps the UX honest for now.
            return gr.update(
                value="✅ You're on the list — we'll reach out when GritAI Studio opens. Thanks!",
                visible=True,
            )

        join.click(_join, inputs=[email, use_case], outputs=waitlist_msg)

    gr.HTML("<div class='gritai-foot'>FLUX.1-schnell &middot; Apache-2.0 &middot; "
            "<a href='https://gritai.solutions'>GRITAI SOLUTIONS</a> &middot; 30-in-15 &#8470;6</div>")

if __name__ == "__main__":
    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))