Update Dockerfile

Dockerfile

@@ -1,89 +1,131 @@
-FROM python:3.10-slim
-
-# System deps (torch + diffusers need these or they sulk)
-RUN apt-get update && apt-get install -y \
-    git \
-    libgl1 \
-    libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/*
-
-# Python deps
-RUN pip install --no-cache-dir \
-    torch \
-    torchvision \
-    torchaudio \
-    diffusers \
-    transformers \
-    accelerate \
-    safetensors \
-    flask \
-    flask-cors \
-    pillow
-
-# Environment
-ENV HOME=/home/sd
-ENV HF_HOME=/home/sd/.cache
-ENV TRANSFORMERS_CACHE=/home/sd/.cache
-ENV DIFFUSERS_CACHE=/home/sd/.cache
-
-# Writable dirs (HF Spaces is picky)
-RUN mkdir -p /home/sd && chmod -R 777 /home/sd
-
-# -------- Flask Stable Diffusion API --------
-RUN cat <<'EOF' > /app.py
-from flask import Flask, request, jsonify, send_file
+FROM ollama/ollama:latest
+
+# ---------------- System + Python ----------------
+RUN apt-get update && apt-get install -y python3 python3-pip && \
+    pip3 install flask flask-cors requests --break-system-packages
+
+# ---------------- Env ----------------
+ENV OLLAMA_HOST=127.0.0.1:11434
+ENV OLLAMA_MODELS=/home/ollama/.ollama/models
+ENV HOME=/home/ollama
+
+# ---------------- Storage ----------------
+RUN mkdir -p /home/ollama/.ollama && chmod -R 777 /home/ollama
+
+# ---------------- Guard API ----------------
+RUN cat <<'EOF' > /guard.py
+from flask import Flask, request, Response, jsonify, stream_with_context
 from flask_cors import CORS
-from diffusers import StableDiffusionPipeline
-import torch
-from io import BytesIO
-import os
+import requests, json, os, datetime, secrets
 
 app = Flask(__name__)
 CORS(app)
 
-MODEL_ID = "runwayml/stable-diffusion-v1-5"
+DB_PATH = "/home/ollama/usage.json"
+WL_PATH = "/home/ollama/whitelist.txt"
+KEY_LIMITS_PATH = "/home/ollama/key_limits.json"
+DEFAULT_LIMIT = 500
+
+# ---------- Init ----------
+os.makedirs("/home/ollama", exist_ok=True)
+
+if not os.path.exists(WL_PATH):
+    open(WL_PATH, "w").close()
+
+def load_limits():
+    if os.path.exists(KEY_LIMITS_PATH):
+        try:
+            return json.load(open(KEY_LIMITS_PATH))
+        except:
+            pass
+    return {}
 
-pipe = StableDiffusionPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+def save_limits(data):
+    json.dump(data, open(KEY_LIMITS_PATH, "w"))
 
+def whitelist():
+    return set(open(WL_PATH).read().split())
+
+# ---------- Health ----------
 @app.route("/", methods=["GET"])
 def health():
-    return "Stable Diffusion API Running", 200
-
+    return "Ollama Guard Running", 200
+
+@app.route("/api/tags", methods=["GET"])
+def tags():
+    try:
+        r = requests.get("http://127.0.0.1:11434/api/tags")
+        return Response(r.content, r.status_code, content_type=r.headers.get("Content-Type"))
+    except:
+        return jsonify({"error": "Ollama starting"}), 503
+
+# ---------- Key Generator ----------
+@app.route("/generate-key", methods=["POST"])
+def gen_key():
+    data = request.get_json() or {}
+    unlimited = data.get("unlimited", False)
+    limit = data.get("limit", DEFAULT_LIMIT)
+
+    key = "sk-" + secrets.token_hex(16)
+
+    with open(WL_PATH, "a") as f:
+        f.write(key + "\n")
+
+    limits = load_limits()
+    limits[key] = "unlimited" if unlimited else int(limit)
+    save_limits(limits)
+
+    return jsonify({
+        "key": key,
+        "limit": limits[key]
+    })
+
+# ---------- Proxy ----------
 @app.route("/api/generate", methods=["POST"])
-def generate():
-    data = request.get_json()
-    prompt = data.get("prompt", "")
-    steps = int(data.get("steps", 25))
-    guidance = float(data.get("guidance", 7.5))
+@app.route("/api/chat", methods=["POST"])
+def proxy():
+    key = request.headers.get("x-api-key", "")
+    if key not in whitelist():
+        return jsonify({"error": "Unauthorized"}), 401
+
+    limits = load_limits()
+    limit = limits.get(key, DEFAULT_LIMIT)
+    unlimited = (limit == "unlimited")
+
+    now = datetime.datetime.now().strftime("%Y-%m")
+    usage = json.load(open(DB_PATH)) if os.path.exists(DB_PATH) else {}
+    used = usage.get(key, {}).get(now, 0)
+
+    if not unlimited and used >= limit:
+        return jsonify({"error": "Monthly limit reached"}), 429
+
+    target = "http://127.0.0.1:11434" + request.path
+    resp = requests.post(target, json=request.json, stream=True, timeout=300)
 
-    if not prompt:
-        return jsonify({"error": "No prompt provided"}), 400
+    if resp.status_code != 200:
+        return jsonify({"error": resp.text}), resp.status_code
 
-    image = pipe(
-        prompt=prompt,
-        num_inference_steps=steps,
-        guidance_scale=guidance
-    ).images[0]
+    usage.setdefault(key, {})[now] = used + 1
+    json.dump(usage, open(DB_PATH, "w"))
 
-    buf = BytesIO()
-    image.save(buf, format="PNG")
-    buf.seek(0)
+    def stream():
+        for c in resp.iter_content(1024):
+            if c: yield c
 
-    return send_file(buf, mimetype="image/png")
+    return Response(stream_with_context(stream()), content_type=resp.headers.get("Content-Type"))
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)
 EOF
 
-# -------- Startup Script --------
+# ---------------- Start Script ----------------
 RUN cat <<'EOF' > /start.sh
 #!/bin/bash
-echo "Starting Stable Diffusion API..."
-python3 /app.py
+ollama serve &
+python3 /guard.py &
+sleep 5
+ollama pull llama3.2
+wait
 EOF
 
 RUN chmod +x /start.sh