Update Dockerfile

Dockerfile

@@ -1,7 +1,8 @@
 # 1. Use Python base
 FROM python:3.10-slim
 
-# 2. Install system dependencies for image processing
+# 2. Install dependencies
+# 'optimum' helps speed up the model on CPU
 RUN apt-get update && apt-get install -y \
     libgl1-mesa-glx libglib2.0-0 procps curl && \
     pip3 install flask flask-cors requests diffusers transformers accelerate torch --extra-index-url https://download.pytorch.org/whl/cpu
@@ -11,7 +12,7 @@ ENV HOME=/home/user
 WORKDIR $HOME
 RUN mkdir -p $HOME/.cache && chmod -R 777 $HOME
 
-# --- 4. CREATE THE GUARD + MODEL SCRIPT (app.py) ---
+# --- 4. THE PYTHON APP (Guard + SD Model) ---
 RUN cat <<EOF > $HOME/app.py
 from flask import Flask, request, Response, jsonify, send_file
 import torch
@@ -22,25 +23,32 @@ import json, os, datetime, io
 app = Flask(__name__)
 CORS(app)
 
+# Note: These files reset every time the Space restarts on the Free Tier
 DB_PATH = "$HOME/usage.json"
-# We read keys from Environment Secrets for easy sharing between Spaces
+WL_PATH = "$HOME/whitelist.txt"
+LIMIT = 500
 UNLIMITED_KEY = "sk-ess4l0ri37"
 
-print("Loading SD Model on CPU... this takes a minute...")
-model_id = "runwayml/stable-diffusion-v1-5"
-# Load specifically for CPU
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+# Seed the whitelist with your unique SD keys
+if not os.path.exists(WL_PATH):
+    with open(WL_PATH, "w") as f:
+        f.write(f"sk-sd-user-1\nsk-sd-pro-99\n{UNLIMITED_KEY}\n")
+
+print("Loading Stable Diffusion on CPU...")
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
 pipe.to("cpu")
-print("Model Loaded!")
+print("Model Ready!")
 
 def get_whitelist():
-    # Pulls from HF Space Secrets 'WL_KEYS' (comma separated)
-    wl_string = os.environ.get("WL_KEYS", UNLIMITED_KEY)
-    return set(k.strip() for k in wl_string.split(","))
+    try:
+        with open(WL_PATH, "r") as f:
+            return set(line.strip() for line in f.readlines())
+    except:
+        return {UNLIMITED_KEY}
 
 @app.route("/", methods=["GET"])
 def health():
-    return "Stable Diffusion CPU API is Running", 200
+    return "Stable Diffusion CPU Proxy is Running", 200
 
 @app.route("/api/generate", methods=["POST"])
 def proxy():
@@ -48,38 +56,59 @@ def proxy():
     
     # 1. Auth Check
     if user_key not in get_whitelist():
-        return jsonify({"error": "Unauthorized: Key not registered"}), 401
+        return jsonify({"error": "Unauthorized: Key not registered for SD"}), 401
 
-    # 2. Process Request
-    data = request.json
-    prompt = data.get("prompt", "a beautiful sunset")
+    # 2. Usage Check
+    is_unlimited = (user_key == UNLIMITED_KEY)
+    if not is_unlimited:
+        now = datetime.datetime.now()
+        month_key = now.strftime("%Y-%m")
+        usage = {}
+        if os.path.exists(DB_PATH):
+            try:
+                with open(DB_PATH, "r") as f:
+                    usage = json.load(f)
+            except: pass
+        
+        key_usage = usage.get(user_key, {}).get(month_key, 0)
+        if key_usage >= LIMIT:
+            return jsonify({"error": f"Monthly limit of {LIMIT} reached"}), 429
 
+    # 3. Image Generation
     try:
-        # num_inference_steps=15 is a good balance for CPU speed
+        data = request.json
+        prompt = data.get("prompt", "a simple drawing of a cat")
+        
+        # CPU speed tweak: Lower steps (15) for faster turnaround
         image = pipe(prompt, num_inference_steps=15).images[0]
         
+        # Log usage
+        if not is_unlimited:
+            if user_key not in usage: usage[user_key] = {}
+            usage[user_key][month_key] = key_usage + 1
+            with open(DB_PATH, "w") as f:
+                json.dump(usage, f)
+
+        # Return Image
         img_io = io.BytesIO()
         image.save(img_io, 'PNG')
         img_io.seek(0)
-        
-        # Add your usage.json logging logic here if needed
-        
         return send_file(img_io, mimetype='image/png')
+        
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        return jsonify({"error": f"Generation Error: {str(e)}"}), 500
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)
 EOF
 
-# --- 5. CREATE THE STARTUP SCRIPT ---
+# --- 5. THE STARTUP SCRIPT ---
 RUN cat <<EOF > $HOME/start.sh
 #!/bin/bash
-echo "Starting Stable Diffusion Guard Service..."
 python3 $HOME/app.py
 EOF
 
 RUN chmod +x $HOME/start.sh
 
-# --- 6. EXECUTION ---
+# --- 6. RUN ---
 ENTRYPOINT ["/bin/bash", "/home/user/start.sh"]