Update Dockerfile

Dockerfile

@@ -1,67 +1,88 @@
-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
+FROM python:3.10-slim
+
+# ---------------- System deps ----------------
+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
 
 # ---------------- Env ----------------
-ENV OLLAMA_HOST=127.0.0.1:11434
-ENV OLLAMA_MODELS=/home/ollama/.ollama/models
-ENV HOME=/home/ollama
+ENV HOME=/home/sd
+ENV HF_HOME=/home/sd/.cache
+ENV TRANSFORMERS_CACHE=/home/sd/.cache
+ENV DIFFUSERS_CACHE=/home/sd/.cache
 
 # ---------------- Storage ----------------
-RUN mkdir -p /home/ollama/.ollama && chmod -R 777 /home/ollama
+RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# ---------------- Guard API ----------------
-RUN cat <<'EOF' > /guard.py
-from flask import Flask, request, Response, jsonify, stream_with_context
+# ---------------- Image Guard API ----------------
+RUN cat <<'EOF' > /app.py
+from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-import requests, json, os, datetime, secrets
+from diffusers import StableDiffusionPipeline
+import torch, os, json, datetime, secrets
+from io import BytesIO
 
 app = Flask(__name__)
 CORS(app)
 
-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 {}
-
-def save_limits(data):
-    json.dump(data, open(KEY_LIMITS_PATH, "w"))
+# -------- Paths --------
+BASE = "/home/sd"
+WL_PATH = f"{BASE}/whitelist.txt"
+USAGE_PATH = f"{BASE}/usage.json"
+LIMITS_PATH = f"{BASE}/limits.json"
 
+DEFAULT_LIMIT = 500
+MODEL_ID = "stable-diffusion-v1-5"
+
+# -------- Init storage --------
+os.makedirs(BASE, exist_ok=True)
+for p in [WL_PATH, USAGE_PATH, LIMITS_PATH]:
+    if not os.path.exists(p):
+        open(p, "w").write("{}" if p.endswith(".json") else "")
+
+# -------- Load model once --------
+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")
+
+# -------- Helpers --------
 def whitelist():
     return set(open(WL_PATH).read().split())
 
-# ---------- Health ----------
-@app.route("/", methods=["GET"])
-def health():
-    return "Ollama Guard Running", 200
-
-@app.route("/api/tags", methods=["GET"])
-def tags():
+def load_json(path):
     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"))
+        return json.load(open(path))
     except:
-        return jsonify({"error": "Ollama starting"}), 503
+        return {}
 
-# ---------- Key Generator ----------
+def save_json(path, data):
+    json.dump(data, open(path, "w"))
+
+# -------- Health --------
+@app.route("/", methods=["GET"])
+def health():
+    return "Image API Running", 200
+
+# -------- Key generator --------
 @app.route("/generate-key", methods=["POST"])
-def gen_key():
+def generate_key():
     data = request.get_json() or {}
     unlimited = data.get("unlimited", False)
     limit = data.get("limit", DEFAULT_LIMIT)
@@ -71,65 +92,68 @@ def gen_key():
     with open(WL_PATH, "a") as f:
         f.write(key + "\n")
 
-    limits = load_limits()
+    limits = load_json(LIMITS_PATH)
     limits[key] = "unlimited" if unlimited else int(limit)
-    save_limits(limits)
+    save_json(LIMITS_PATH, limits)
 
     return jsonify({
         "key": key,
         "limit": limits[key]
     })
 
-# ---------- Proxy ----------
+# -------- Image generation --------
 @app.route("/api/generate", methods=["POST"])
-@app.route("/api/chat", methods=["POST"])
-def proxy():
+def generate():
     key = request.headers.get("x-api-key", "")
     if key not in whitelist():
         return jsonify({"error": "Unauthorized"}), 401
 
-    limits = load_limits()
+    data = request.get_json() or {}
+    prompt = data.get("prompt", "").strip()
+    steps = int(data.get("steps", 25))
+    guidance = float(data.get("guidance", 7.5))
+
+    if not prompt:
+        return jsonify({"error": "Prompt required"}), 400
+
+    limits = load_json(LIMITS_PATH)
+    usage = load_json(USAGE_PATH)
+
     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)
+    month = datetime.datetime.now().strftime("%Y-%m")
+    used = usage.get(key, {}).get(month, 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)
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance
+    ).images[0]
 
-    if resp.status_code != 200:
-        return jsonify({"error": resp.text}), resp.status_code
+    usage.setdefault(key, {})[month] = used + 1
+    save_json(USAGE_PATH, usage)
 
-    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 Response(stream_with_context(stream()), content_type=resp.headers.get("Content-Type"))
+    return send_file(buf, mimetype="image/png")
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)
 EOF
 
-# ---------------- Start Script ----------------
+# ---------------- Start ----------------
 RUN cat <<'EOF' > /start.sh
 #!/bin/bash
-ollama serve &
-python3 /guard.py &
-sleep 5
-ollama pull llama3.2
-wait
+python3 /app.py
 EOF
 
 RUN chmod +x /start.sh
 
 EXPOSE 7860
-
 ENTRYPOINT ["/bin/bash", "/start.sh"]