Spaces:
Runtime error
Runtime error
UI: GritAI engineered/industrial design system (dark, grit-orange, Archivo+IBM Plex)
d3c0565 verified | #!/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 | |
| 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) | |
| 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">▸ 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 · $0/image · 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 · Apache-2.0 · " | |
| "<a href='https://gritai.solutions'>GRITAI SOLUTIONS</a> · 30-in-15 №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))) | |