Spaces:
Paused
Paused
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import io | |
| import json | |
| import os | |
| import time | |
| import secrets | |
| import datetime | |
| from fastapi import Request | |
| from fastapi.responses import JSONResponse, Response | |
| # ── ZeroGPU / SD1.5 auto-detect ────────────────────────────────────────────── | |
| HAS_ZERO_GPU = False | |
| HAS_CUDA = False | |
| try: | |
| import spaces | |
| # Only enable if actually running on ZeroGPU hardware | |
| import os | |
| if os.environ.get("SPACE_HARDWARE") in ("gpu-t4-s", "gpu-t4-m", "gpu-a10g-s", "gpu-a10g-l", "zero-gpu"): | |
| HAS_ZERO_GPU = True | |
| print("✅ ZeroGPU detected — SD1.5 'best' mode will be available") | |
| else: | |
| print("ℹ️ spaces imported but no ZeroGPU hardware — LaMa-only mode") | |
| except ImportError: | |
| print("ℹ️ No ZeroGPU — LaMa-only mode (fast + quality)") | |
| try: | |
| import torch | |
| HAS_CUDA = torch.cuda.is_available() | |
| except: | |
| pass | |
| # ── HF Secret auto-update (like FreeBG) ────────────────────────────────────── | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") | |
| HF_REPO_ID = os.environ.get("HF_REPO_ID", "") # e.g. "freebg/magic-eraser" | |
| def save_keys_to_hf(keys: dict) -> bool: | |
| """Push updated API_KEYS_JSON straight to HF Space Secret — Space restarts automatically.""" | |
| if not HF_TOKEN or not HF_REPO_ID: | |
| return False | |
| try: | |
| from huggingface_hub import HfApi | |
| api = HfApi(token=HF_TOKEN) | |
| api.add_space_secret( | |
| repo_id=HF_REPO_ID, | |
| key="API_KEYS_JSON", | |
| value=json.dumps(keys), | |
| ) | |
| return True | |
| except Exception as e: | |
| print(f"⚠️ HF Secret update failed: {e}") | |
| return False | |
| # ── LaMa model ──────────────────────────────────────────────────────────────── | |
| lama_model = None | |
| def load_lama(): | |
| global lama_model | |
| if lama_model is not None: | |
| return lama_model | |
| try: | |
| from simple_lama_inpainting import SimpleLama | |
| lama_model = SimpleLama() | |
| print("✅ LaMa loaded") | |
| return lama_model | |
| except Exception as e: | |
| print(f"❌ LaMa load error: {e}") | |
| return None | |
| # ── SD1.5 pipeline (ZeroGPU / CUDA only) ───────────────────────────────────── | |
| sd_pipe = None | |
| def load_sd(): | |
| return None | |
| # ── Key / Plan config ───────────────────────────────────────────────────────── | |
| MASTER_KEY = os.environ.get("MASTER_API_KEY", "freebg-master-change-this") | |
| ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "admin123") | |
| PLAN_LIMITS = { | |
| "free": {"daily": 10, "models": ["fast"]}, | |
| "starter": {"daily": 100, "models": ["fast", "quality"]}, | |
| "pro": {"daily": 500, "models": ["fast", "quality", "best"]}, | |
| "master": {"daily": 99999, "models": ["fast", "quality", "best"]}, | |
| } | |
| # In-memory keys cache (reloaded from env each call so secrets take effect after restart) | |
| usage_cache: dict = {} | |
| def get_today() -> str: | |
| return datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d") | |
| def load_keys() -> dict: | |
| raw = os.environ.get("API_KEYS_JSON", "{}") | |
| try: | |
| return json.loads(raw) | |
| except: | |
| return {} | |
| def check_and_count(api_key: str, model: str) -> tuple[bool, str]: | |
| if api_key == MASTER_KEY: | |
| return True, "" | |
| keys = load_keys() | |
| if api_key not in keys: | |
| return False, "Invalid API key. Get one at freebg.site" | |
| info = keys[api_key] | |
| plan = info.get("plan", "free") | |
| lims = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) | |
| if model not in lims["models"]: | |
| return False, f"Model '{model}' not available on {plan} plan. Upgrade at freebg.site" | |
| today = get_today() | |
| cache_key = f"{api_key}:{today}" | |
| used = usage_cache.get(cache_key, 0) | |
| if used >= lims["daily"]: | |
| return False, f"Daily limit reached ({lims['daily']}/day on {plan} plan)" | |
| usage_cache[cache_key] = used + 1 | |
| return True, "" | |
| def get_usage_info(api_key: str) -> dict: | |
| """Always returns a dict with 'models' key — no KeyError possible.""" | |
| plan_info = PLAN_LIMITS.get("master") | |
| if api_key == MASTER_KEY: | |
| today = get_today() | |
| used = usage_cache.get(f"{api_key}:{today}", 0) | |
| return { | |
| "plan": "master", | |
| "daily_limit": plan_info["daily"], | |
| "used_today": used, | |
| "remaining": plan_info["daily"] - used, | |
| "models": plan_info["models"], | |
| "owner": "freebg.site", | |
| } | |
| keys = load_keys() | |
| if api_key not in keys: | |
| return {"error": "Invalid API key"} | |
| info = keys[api_key] | |
| plan = info.get("plan", "free") | |
| lims = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"]) | |
| today = get_today() | |
| used = usage_cache.get(f"{api_key}:{today}", 0) | |
| return { | |
| "plan": plan, | |
| "daily_limit": lims["daily"], | |
| "used_today": used, | |
| "remaining": max(0, lims["daily"] - used), | |
| "models": lims["models"], | |
| "owner": info.get("owner", ""), | |
| } | |
| # ── Inpainting ──────────────────────────────────────────────────────────────── | |
| def _prep_mask(mask: Image.Image) -> Image.Image: | |
| arr = np.array(mask.convert("L")) | |
| return Image.fromarray(((arr > 127).astype(np.uint8) * 255), "L") | |
| def run_lama(image: Image.Image, mask: Image.Image) -> Image.Image: | |
| lama = load_lama() | |
| if lama is None: | |
| raise RuntimeError("LaMa model failed to load") | |
| return lama(image.convert("RGB"), _prep_mask(mask)) | |
| def run_lama_quality(image: Image.Image, mask: Image.Image) -> Image.Image: | |
| from PIL import ImageFilter | |
| return run_lama(image, mask).filter(ImageFilter.UnsharpMask(radius=1.5, percent=120, threshold=3)) | |
| def run_best(image: Image.Image, mask: Image.Image) -> Image.Image: | |
| return run_lama_quality(image, mask) | |
| if HAS_ZERO_GPU: | |
| import spaces as _spaces | |
| run_lama = _spaces.GPU(run_lama) | |
| run_lama_quality = _spaces.GPU(run_lama_quality) | |
| run_best = _spaces.GPU(run_best) | |
| # ── Gradio helpers ──────────────────────────────────────────────────────────── | |
| def erase_ui(composite, model: str, api_key: str): | |
| if composite is None: | |
| return None, "❌ Please upload and paint an image first" | |
| background = composite.get("background") | |
| layers = composite.get("layers", []) | |
| if background is None: | |
| return None, "❌ No image provided" | |
| image = background.convert("RGB") | |
| if not layers: | |
| return None, "❌ Please paint over the area to erase" | |
| r, g, b, a = layers[0].convert("RGBA").split() | |
| mask = a | |
| if np.array(mask).max() == 0: | |
| return None, "❌ No brush strokes found — paint over the area to erase" | |
| key = (api_key or "").strip() or MASTER_KEY | |
| allowed, err = check_and_count(key, model.lower()) | |
| if not allowed: | |
| return None, f"❌ {err}" | |
| try: | |
| t0 = time.time() | |
| fn = {"fast": run_lama, "quality": run_lama_quality, "best": run_best}.get(model, run_lama) | |
| result = fn(image, mask) | |
| return result, f"✅ Done in {time.time()-t0:.1f}s | Model: {model}" | |
| except Exception as e: | |
| return None, f"❌ Error: {e}" | |
| def check_usage_ui(api_key: str) -> str: | |
| key = (api_key or "").strip() | |
| if not key: | |
| return "❌ Enter your API key above" | |
| info = get_usage_info(key) | |
| if "error" in info: | |
| return f"❌ {info['error']}" | |
| return ( | |
| f"**Plan:** {info['plan'].upper()}\n\n" | |
| f"**Calls today:** {info['used_today']} / {info['daily_limit']}\n\n" | |
| f"**Remaining:** {info['remaining']}\n\n" | |
| f"**Models:** {', '.join(info['models'])}\n\n" | |
| f"**Owner:** {info['owner']}" | |
| ) | |
| # ── Admin helpers ───────────────────────────────────────────────────────────── | |
| def _auth(pw: str) -> bool: | |
| return pw == ADMIN_PASSWORD | |
| def admin_add_customer(pw, email, plan, custom_limit): | |
| if not _auth(pw): | |
| return "❌ Wrong password", "", "" | |
| if plan not in PLAN_LIMITS: | |
| return f"❌ Invalid plan", "", "" | |
| prefix = {"free":"me-fr","starter":"me-st","pro":"me-pr","master":"me-ma"}[plan] | |
| key = f"freebg-{prefix}-{secrets.token_hex(4)}-{secrets.token_hex(4)}" | |
| keys = load_keys() | |
| entry = {"owner": email, "plan": plan, "created": get_today(), "active": True} | |
| if custom_limit: | |
| try: | |
| entry["custom_daily"] = int(custom_limit) | |
| except: | |
| pass | |
| keys[key] = entry | |
| saved = save_keys_to_hf(keys) | |
| hf_note = "✅ Auto-saved to HF Secret. Space restarts in ~30s." if saved else "⚠️ HF auto-save failed — paste JSON manually." | |
| return f"✅ Key created for {email} ({plan})\n\nKey: **{key}**\n\n{hf_note}", key, json.dumps(keys, indent=2) | |
| def admin_get_customers(pw): | |
| if not _auth(pw): | |
| return "❌ Wrong password" | |
| keys = load_keys() | |
| today = get_today() | |
| rows = [] | |
| for k, v in keys.items(): | |
| used = usage_cache.get(f"{k}:{today}", 0) | |
| plan = v.get("plan","free") | |
| limit = v.get("custom_daily", PLAN_LIMITS.get(plan,{}).get("daily","?")) | |
| rows.append(f"| `{k}` | {v.get('owner','')} | {plan} | {used} | {limit} | {v.get('created','')} |") | |
| if not rows: | |
| return "No customers yet." | |
| header = "| API Key | Owner | Plan | Calls Today | Daily Limit | Created |\n|---|---|---|---|---|---|" | |
| return header + "\n" + "\n".join(rows) | |
| def admin_upgrade(pw, key, new_plan): | |
| if not _auth(pw): | |
| return "❌ Wrong password", "" | |
| keys = load_keys() | |
| if key not in keys: | |
| return "❌ Key not found", "" | |
| old = keys[key].get("plan","?") | |
| keys[key]["plan"] = new_plan | |
| saved = save_keys_to_hf(keys) | |
| note = "✅ Auto-saved to HF Secret." if saved else "⚠️ Paste JSON manually." | |
| return f"✅ {key}\n{old} → {new_plan}\n{note}", json.dumps(keys, indent=2) | |
| def admin_delete(pw, key): | |
| if not _auth(pw): | |
| return "❌ Wrong password", "" | |
| keys = load_keys() | |
| if key not in keys: | |
| return "❌ Key not found", "" | |
| del keys[key] | |
| saved = save_keys_to_hf(keys) | |
| note = "✅ Auto-saved to HF Secret." if saved else "⚠️ Paste JSON manually." | |
| return f"✅ Deleted: {key}\n{note}", json.dumps(keys, indent=2) | |
| def admin_stats(pw): | |
| if not _auth(pw): | |
| return "❌ Wrong password" | |
| keys = load_keys() | |
| today = get_today() | |
| plans = {"free":0,"starter":0,"pro":0} | |
| calls = 0 | |
| for k, v in keys.items(): | |
| p = v.get("plan","free") | |
| if p in plans: | |
| plans[p] += 1 | |
| calls += usage_cache.get(f"{k}:{today}", 0) | |
| rev = plans["starter"]*9 + plans["pro"]*29 | |
| return ( | |
| f"👥 **Total Customers:** {len(keys)}\n\n" | |
| f"🆓 Free: {plans['free']}\n\n" | |
| f"⭐ Starter ($9): {plans['starter']}\n\n" | |
| f"🏆 Pro ($29): {plans['pro']}\n\n" | |
| f"📊 **Total Calls Today:** {calls}\n\n" | |
| f"💰 **Est. Monthly Revenue:** ${rev}" | |
| ) | |
| def admin_export_json(pw): | |
| if not _auth(pw): | |
| return "❌ Wrong password" | |
| return json.dumps(load_keys(), indent=2) | |
| # ── Build Gradio UI ─────────────────────────────────────────────────────────── | |
| gpu_note = "⚡ ZeroGPU active — best mode uses SD1.5 inpainting" if HAS_ZERO_GPU else "🖥️ CPU mode — fast & quality use LaMa. Best = LaMa fallback. Enable ZeroGPU for SD1.5." | |
| with gr.Blocks(title="🪄 FreeME — Magic Eraser API") as demo: | |
| gr.HTML(""" | |
| <div style="text-align:center;padding:20px 0 8px;"> | |
| <h1 style="font-size:1.8rem;font-weight:800;margin:0;">🪄 FreeME — Magic Eraser API</h1> | |
| <p style="color:#94a3b8;margin:6px 0 0;">3 Models · API Key Auth · <a href="https://freebg.site" style="color:#f97316;">freebg.site</a></p> | |
| </div>""") | |
| with gr.Tabs(): | |
| # ── Tab 1: Magic Erase ──────────────────────────────────────────────── | |
| with gr.TabItem("🪄 Magic Erase"): | |
| gr.Markdown(f"> {gpu_note}") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| api_key_in = gr.Textbox(label="🔑 API Key", placeholder="freebg-me-xxxx", type="password") | |
| model_in = gr.Radio(["fast","quality","best"], value="fast", label="⚙️ Model", | |
| info="fast=LaMa | quality=LaMa+sharpen | best=SD1.5 on GPU / LaMa fallback on CPU") | |
| img_editor = gr.ImageEditor( | |
| label="🖌️ Upload image then paint over what to erase", | |
| type="pil", | |
| brush=gr.Brush(colors=["#FF0000"], default_size=20), | |
| height=440, | |
| ) | |
| erase_btn = gr.Button("🪄 Erase Now", variant="primary") | |
| with gr.Column(scale=1): | |
| result_out = gr.Image(label="✨ Result", type="pil", height=360) | |
| status_out = gr.Textbox(label="Status", interactive=False) | |
| erase_btn.click(fn=erase_ui, inputs=[img_editor, model_in, api_key_in], outputs=[result_out, status_out]) | |
| gr.HTML(""" | |
| <div style="margin-top:20px;padding:14px 18px;background:#1e293b;border-radius:8px;"> | |
| <b>📦 Plans</b> | |
| <table style="width:100%;margin-top:8px;border-collapse:collapse;"> | |
| <tr style="border-bottom:1px solid #334155;"><th style="text-align:left;padding:6px;">Plan</th><th>Daily</th><th>Models</th><th>Price</th></tr> | |
| <tr><td style="padding:5px;">Free</td><td style="text-align:center;">10</td><td style="text-align:center;">fast</td><td style="text-align:center;">$0</td></tr> | |
| <tr><td style="padding:5px;">Starter</td><td style="text-align:center;">100</td><td style="text-align:center;">fast + quality</td><td style="text-align:center;">$9/mo</td></tr> | |
| <tr><td style="padding:5px;">Pro</td><td style="text-align:center;">500</td><td style="text-align:center;">all 3</td><td style="text-align:center;">$29/mo</td></tr> | |
| </table> | |
| <p style="margin:8px 0 0;color:#94a3b8;font-size:.85rem;">📧 <a href="mailto:support@freebg.site" style="color:#f97316;">support@freebg.site</a></p> | |
| </div>""") | |
| # ── Tab 2: My Usage ─────────────────────────────────────────────────── | |
| with gr.TabItem("📊 My Usage"): | |
| gr.Markdown("### Check Your API Key Usage") | |
| with gr.Row(): | |
| usage_key_in = gr.Textbox(label="🔑 Your API Key", type="password", | |
| placeholder="freebg-me-xxxx", scale=4) | |
| usage_btn = gr.Button("Check Usage", scale=1) | |
| usage_out = gr.Markdown() | |
| usage_btn.click(fn=check_usage_ui, inputs=usage_key_in, outputs=usage_out) | |
| # ── Tab 3: API Docs ─────────────────────────────────────────────────── | |
| with gr.TabItem("📖 API Docs"): | |
| gr.Markdown(f""" | |
| ## FreeME Magic Eraser — REST API | |
| **Base URL:** `https://freebg-magic-eraser.hf.space` | |
| **GPU Status:** `{"ZeroGPU Active ⚡" if HAS_ZERO_GPU else "CPU Mode 🖥️"}` | |
| --- | |
| ### `POST /api/magic-erase` | |
| | Field | Type | Required | Description | | |
| |-------|------|----------|-------------| | |
| | `image` | file | ✅ | Input image JPG/PNG | | |
| | `mask` | file | ✅ | Mask PNG — white=erase, black=keep | | |
| | `model` | string | ❌ | `fast` / `quality` / `best` (default: fast) | | |
| **Header:** `X-API-Key: your-key` | |
| **Response:** `image/png` | |
| --- | |
| ### `GET /api/usage` | |
| **Header:** `X-API-Key: your-key` | |
| ```json | |
| {{ | |
| "plan": "starter", | |
| "daily_limit": 100, | |
| "used_today": 23, | |
| "remaining": 77, | |
| "models": ["fast", "quality"], | |
| "owner": "user@email.com" | |
| }} | |
| ``` | |
| --- | |
| ### Python | |
| ```python | |
| import requests | |
| r = requests.post( | |
| "https://freebg-magic-eraser.hf.space/api/magic-erase", | |
| headers={{"X-API-Key": "your-key"}}, | |
| files={{"image": open("photo.jpg","rb"), "mask": open("mask.png","rb")}}, | |
| data={{"model": "quality"}} | |
| ) | |
| open("result.png","wb").write(r.content) | |
| ``` | |
| ### JavaScript | |
| ```javascript | |
| const fd = new FormData(); | |
| fd.append("image", imageFile); | |
| fd.append("mask", maskBlob, "mask.png"); | |
| fd.append("model", "fast"); | |
| const res = await fetch("https://freebg-magic-eraser.hf.space/api/magic-erase", {{ | |
| method: "POST", headers: {{"X-API-Key": "your-key"}}, body: fd | |
| }}); | |
| const url = URL.createObjectURL(await res.blob()); | |
| ``` | |
| """) | |
| # ── Tab 4: Admin ────────────────────────────────────────────────────── | |
| with gr.TabItem("🔐 Admin"): | |
| gr.HTML("<h3 style='margin:12px 0 4px;'>👑 FreeME Admin Panel</h3>") | |
| with gr.Group(visible=True) as login_group: | |
| admin_pw_login = gr.Textbox(label="Admin Password", type="password", placeholder="Enter admin password") | |
| login_btn = gr.Button("Login", variant="primary") | |
| login_err = gr.Markdown(visible=False) | |
| with gr.Group(visible=False) as panel_group: | |
| admin_pw = gr.State("") # store password after login | |
| with gr.Tabs(): | |
| # Add Customer | |
| with gr.TabItem("➕ Add Customer"): | |
| gr.Markdown("Fill details → Generate Key → Key is **automatically saved** to HF Secret") | |
| with gr.Row(): | |
| ac_email = gr.Textbox(label="Customer Email", placeholder="customer@gmail.com", scale=3) | |
| ac_plan = gr.Dropdown(["free","starter","pro"], value="starter", label="Plan", scale=2) | |
| ac_limit = gr.Textbox(label="Custom Daily Limit (optional)", placeholder="leave empty = plan default", scale=2) | |
| ac_btn = gr.Button("🔑 Generate API Key + Save to HF Secret", variant="primary") | |
| ac_status = gr.Markdown() | |
| ac_key = gr.Textbox(label="🔑 Generated Key (copy & send to customer)", interactive=False) | |
| ac_json = gr.Code(label="Updated JSON (auto-saved — for manual backup)", language="json") | |
| ac_btn.click(fn=admin_add_customer, | |
| inputs=[admin_pw, ac_email, ac_plan, ac_limit], | |
| outputs=[ac_status, ac_key, ac_json]) | |
| # Customers list | |
| with gr.TabItem("👥 Customers"): | |
| cust_refresh = gr.Button("🔄 Refresh List") | |
| cust_out = gr.Markdown() | |
| with gr.Row(): | |
| upg_key = gr.Textbox(label="API Key to Edit", placeholder="freebg-me-xx-xxxxxxxx-xxxxxxxx", scale=3) | |
| upg_plan = gr.Dropdown(["free","starter","pro","master"], value="pro", label="New Plan", scale=2) | |
| upg_btn = gr.Button("⬆️ Upgrade Plan + Save to HF Secret") | |
| upg_status = gr.Textbox(label="Status", interactive=False) | |
| upg_json = gr.Code(label="Updated JSON", language="json") | |
| gr.Markdown("---") | |
| del_key = gr.Textbox(label="API Key to Delete", placeholder="freebg-me-xx-xxxxxxxx-xxxxxxxx") | |
| del_btn = gr.Button("🗑️ Delete Key + Save to HF Secret", variant="stop") | |
| del_status = gr.Textbox(label="Status", interactive=False) | |
| del_json = gr.Code(label="Updated JSON", language="json") | |
| cust_refresh.click(fn=admin_get_customers, inputs=admin_pw, outputs=cust_out) | |
| upg_btn.click(fn=admin_upgrade, inputs=[admin_pw, upg_key, upg_plan], outputs=[upg_status, upg_json]) | |
| del_btn.click(fn=admin_delete, inputs=[admin_pw, del_key], outputs=[del_status, del_json]) | |
| # Stats | |
| with gr.TabItem("📊 Stats"): | |
| stats_btn = gr.Button("🔄 Refresh Stats") | |
| stats_out = gr.Markdown() | |
| stats_btn.click(fn=admin_stats, inputs=admin_pw, outputs=stats_out) | |
| # Export JSON | |
| with gr.TabItem("📋 Export JSON"): | |
| gr.Markdown("Use this if you want to manually paste into HF Secret. **Auto-update already happens** when you Add/Edit/Delete customers.") | |
| exp_btn = gr.Button("📋 Get Current JSON") | |
| exp_out = gr.Code(language="json") | |
| exp_btn.click(fn=admin_export_json, inputs=admin_pw, outputs=exp_out) | |
| def do_login(pw): | |
| if pw == ADMIN_PASSWORD: | |
| return gr.update(visible=False), gr.update(visible=True), pw, gr.update(visible=False) | |
| return gr.update(visible=True), gr.update(visible=False), "", gr.update(visible=True, value="❌ Wrong password") | |
| login_btn.click( | |
| fn=do_login, | |
| inputs=admin_pw_login, | |
| outputs=[login_group, panel_group, admin_pw, login_err] | |
| ) | |
| # ── FastAPI REST endpoints ──────────────────────────────────────────────────── | |
| app = demo.app | |
| async def api_magic_erase(request: Request): | |
| import traceback | |
| t0 = time.time() | |
| try: | |
| form = await request.form() | |
| api_key = request.headers.get("X-API-Key", "") | |
| model = form.get("model", "fast").lower() | |
| if model not in ["fast","quality","best"]: | |
| model = "fast" | |
| allowed, err = check_and_count(api_key, model) | |
| if not allowed: | |
| return JSONResponse({"success": False, "error": err}, status_code=403) | |
| img_file = form.get("image") | |
| mask_file = form.get("mask") | |
| if img_file is None or mask_file is None: | |
| return JSONResponse({"success": False, "error": "Both 'image' and 'mask' files required"}, status_code=400) | |
| image = Image.open(io.BytesIO(await img_file.read())).convert("RGB") | |
| mask = Image.open(io.BytesIO(await mask_file.read())).convert("L") | |
| if mask.size != image.size: | |
| mask = mask.resize(image.size, Image.NEAREST) | |
| fn = {"fast": run_lama, "quality": run_lama_quality, "best": run_best}[model] | |
| result = fn(image, mask) | |
| buf = io.BytesIO() | |
| result.save(buf, format="PNG") | |
| return Response( | |
| content=buf.getvalue(), | |
| media_type="image/png", | |
| headers={ | |
| "X-Model": model, | |
| "X-Processing-Time": f"{time.time()-t0:.2f}s", | |
| "X-GPU": "zerogpu" if HAS_ZERO_GPU else ("cuda" if HAS_CUDA else "cpu"), | |
| } | |
| ) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return JSONResponse({"success": False, "error": str(e)}, status_code=500) | |
| async def api_usage(request: Request): | |
| return JSONResponse(get_usage_info(request.headers.get("X-API-Key",""))) | |
| async def api_health(): | |
| return JSONResponse({ | |
| "status": "ok", | |
| "lama": lama_model is not None, | |
| "zero_gpu": HAS_ZERO_GPU, | |
| "cuda": HAS_CUDA, | |
| "mode": "zerogpu" if HAS_ZERO_GPU else ("cuda" if HAS_CUDA else "cpu"), | |
| }) | |
| # ── Launch ──────────────────────────────────────────────────────────────────── | |
| load_lama() | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False) |