#!/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('''
GRITAI
AI Image Generator
▸ LiveSDVOSBLawton OK

Type an idea. Get an image in seconds.

FLUX.1-schnell on free GPU · $0/image · Project #6 of 30 in 15 days.

''') 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( "
\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" "
" ) 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("
FLUX.1-schnell · Apache-2.0 · " "GRITAI SOLUTIONS · 30-in-15 №6
") if __name__ == "__main__": demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))