Spaces:
Running
Running
| import os | |
| import sys | |
| import threading | |
| import warnings | |
| import io as _io | |
| warnings.filterwarnings("ignore") | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
| os.environ["OMP_NUM_THREADS"] = "2" | |
| os.environ["ORT_LOGGING_LEVEL"] = "3" | |
| # Suppress GPU stderr on import | |
| _old_stderr = sys.stderr | |
| sys.stderr = _io.StringIO() | |
| try: | |
| import onnxruntime as ort | |
| ort.set_default_logger_severity(3) | |
| finally: | |
| sys.stderr = _old_stderr | |
| from flask import Flask, request, render_template, send_file, jsonify | |
| from flask_cors import CORS | |
| from rembg import remove, new_session | |
| from PIL import Image, ImageFilter | |
| import io | |
| import time | |
| app = Flask(__name__) | |
| CORS(app) | |
| # Only ONE model loaded at a time to stay within 512MB RAM | |
| MODELS = { | |
| "general": "u2net", | |
| "portrait": "u2net_human_seg", | |
| "anime": "isnet-anime", | |
| "product": "silueta", # smallest at 43MB | |
| } | |
| SESS_OPTS = ort.SessionOptions() | |
| SESS_OPTS.inter_op_num_threads = 1 | |
| SESS_OPTS.intra_op_num_threads = 1 | |
| SESS_OPTS.log_severity_level = 3 | |
| CPU_PROVIDERS = ["CPUExecutionProvider"] | |
| # Only keep ONE session at a time to save memory | |
| _current_session = None | |
| _current_model = None | |
| _lock = threading.Lock() | |
| def get_session(model_name): | |
| global _current_session, _current_model | |
| if _current_model == model_name and _current_session is not None: | |
| return _current_session | |
| with _lock: | |
| if _current_model != model_name: | |
| print(f"[CutOut] Switching model: {_current_model} → {model_name}", flush=True) | |
| # Free old session from memory | |
| _current_session = None | |
| _current_model = None | |
| import gc | |
| gc.collect() | |
| try: | |
| _current_session = new_session( | |
| model_name, | |
| sess_options=SESS_OPTS, | |
| providers=CPU_PROVIDERS, | |
| ) | |
| _current_model = model_name | |
| print(f"[CutOut] Ready: {model_name}", flush=True) | |
| except Exception as e: | |
| print(f"[CutOut] Failed {model_name}: {e}", flush=True) | |
| raise e | |
| return _current_session | |
| def preload_default(): | |
| """Only preload the general model at startup — 176MB fits fine.""" | |
| time.sleep(2) | |
| try: | |
| get_session("u2net") | |
| print("[CutOut] Default model ready!", flush=True) | |
| except Exception as e: | |
| print(f"[CutOut] Preload failed: {e}", flush=True) | |
| threading.Thread(target=preload_default, daemon=True).start() | |
| def index(): | |
| return render_template("index.html") | |
| def ping(): | |
| return "pong", 200 | |
| def health(): | |
| return jsonify({ | |
| "status": "ok", | |
| "models": list(MODELS.keys()), | |
| "loaded": _current_model, | |
| "version": "2.6", | |
| }) | |
| def remove_background(): | |
| start = time.time() | |
| if "image" not in request.files: | |
| return jsonify({"error": "No file uploaded."}), 400 | |
| file = request.files["image"] | |
| if not file.filename: | |
| return jsonify({"error": "No file selected."}), 400 | |
| model_key = request.form.get("model", "general") | |
| bg_color = request.form.get("bg_color", "").strip() | |
| feather = int(request.form.get("feather", 0)) | |
| shadow = request.form.get("shadow", "false") == "true" | |
| shadow_blur = int(request.form.get("shadow_blur", 12)) | |
| shadow_opac = int(request.form.get("shadow_opacity", 60)) | |
| out_format = request.form.get("format", "png").lower() | |
| scale = float(request.form.get("scale", 1.0)) | |
| model_name = MODELS.get(model_key, "u2net") | |
| try: | |
| input_bytes = file.read() | |
| session = get_session(model_name) | |
| output_bytes = remove(input_bytes, session=session) | |
| except Exception as e: | |
| print(f"[CutOut] Error: {e}", flush=True) | |
| return jsonify({"error": str(e)}), 500 | |
| img = Image.open(io.BytesIO(output_bytes)).convert("RGBA") | |
| if feather > 0: | |
| r, g, b, a = img.split() | |
| a = a.filter(ImageFilter.GaussianBlur(radius=feather)) | |
| img = Image.merge("RGBA", (r, g, b, a)) | |
| if scale != 1.0: | |
| img = img.resize( | |
| (int(img.width * scale), int(img.height * scale)), | |
| Image.LANCZOS | |
| ) | |
| if bg_color or shadow: | |
| canvas = Image.new("RGBA", img.size, (0, 0, 0, 0)) | |
| if shadow: | |
| _, _, _, mask = img.split() | |
| sh = Image.new("RGBA", img.size, (0, 0, 0, int(shadow_opac * 2.55))) | |
| sh.putalpha(mask) | |
| sh = sh.filter(ImageFilter.GaussianBlur(radius=shadow_blur)) | |
| sc = Image.new("RGBA", img.size, (0, 0, 0, 0)) | |
| sc.paste(sh, (6, 8)) | |
| canvas = Image.alpha_composite(canvas, sc) | |
| if bg_color.startswith("#") and len(bg_color) >= 7: | |
| try: | |
| h = bg_color.lstrip("#") | |
| bl = Image.new("RGBA", img.size, ( | |
| int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16), 255 | |
| )) | |
| canvas = Image.alpha_composite(bl, canvas) | |
| except ValueError: | |
| pass | |
| img = Image.alpha_composite(canvas, img) | |
| buf = io.BytesIO() | |
| elapsed = round((time.time() - start) * 1000) | |
| if out_format == "jpg": | |
| flat = Image.new("RGB", img.size, (255, 255, 255)) | |
| flat.paste(img, mask=img.split()[3]) | |
| flat.save(buf, "JPEG", quality=95) | |
| mime, fname = "image/jpeg", "cutout.jpg" | |
| elif out_format == "webp": | |
| img.save(buf, "WEBP", quality=95) | |
| mime, fname = "image/webp", "cutout.webp" | |
| else: | |
| img.save(buf, "PNG") | |
| mime, fname = "image/png", "cutout.png" | |
| buf.seek(0) | |
| print(f"[CutOut] {elapsed}ms {model_name} {img.width}x{img.height}", flush=True) | |
| resp = send_file(buf, mimetype=mime, as_attachment=False, download_name=fname) | |
| resp.headers["X-Processing-Time"] = str(elapsed) | |
| resp.headers["X-Image-Width"] = str(img.width) | |
| resp.headers["X-Image-Height"] = str(img.height) | |
| return resp | |
| if __name__ == "__main__": | |
| port = int(os.environ.get("PORT", 5000)) | |
| print(f"[CutOut] Starting on port {port}", flush=True) | |
| app.run(host="0.0.0.0", port=port, debug=False) | |