Update Dockerfile

Dockerfile

@@ -1,114 +1,58 @@
-# 1. Use Python base
+# 1. Use a more stable Python image
 FROM python:3.10-slim
 
-# 2. Install dependencies
-# 'optimum' helps speed up the model on CPU
+# 2. System dependencies
 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
+    rm -rf /var/lib/apt/lists/*
 
-# 3. Set up environment
+# 3. Python dependencies
+RUN pip3 install --no-cache-dir \
+    flask flask-cors requests diffusers transformers accelerate torch \
+    --extra-index-url https://download.pytorch.org/whl/cpu
+
+# 4. Set up non-root user (Required for HF)
+RUN useradd -m -u 1000 user
+USER user
 ENV HOME=/home/user
+ENV PATH=/home/user/.local/bin:$PATH
 WORKDIR $HOME
-RUN mkdir -p $HOME/.cache && chmod -R 777 $HOME
-
-# --- 4. THE PYTHON APP (Guard + SD Model) ---
-RUN cat <<EOF > $HOME/app.py
-from flask import Flask, request, Response, jsonify, send_file
-import torch
-from diffusers import StableDiffusionPipeline
-from flask_cors import CORS
-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"
-WL_PATH = "$HOME/whitelist.txt"
-LIMIT = 500
-UNLIMITED_KEY = "sk-ess4l0ri37"
-
-# 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 Ready!")
-
-def get_whitelist():
-    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 Proxy is Running", 200
-
-@app.route("/api/generate", methods=["POST"])
-def proxy():
-    user_key = request.headers.get("x-api-key", "")
-    
-    # 1. Auth Check
-    if user_key not in get_whitelist():
-        return jsonify({"error": "Unauthorized: Key not registered for SD"}), 401
-
-    # 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:
-        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)
-        return send_file(img_io, mimetype='image/png')
-        
-    except Exception as e:
-        return jsonify({"error": f"Generation Error: {str(e)}"}), 500
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
-EOF
-
-# --- 5. THE STARTUP SCRIPT ---
-RUN cat <<EOF > $HOME/start.sh
-#!/bin/bash
-python3 $HOME/app.py
-EOF
-
-RUN chmod +x $HOME/start.sh
 
-# --- 6. RUN ---
-ENTRYPOINT ["/bin/bash", "/home/user/start.sh"]
+# 5. Create the app file directly via a more robust method
+# Using printf instead of cat <<EOF to avoid shell/indentation errors
+RUN printf 'from flask import Flask, request, jsonify, send_file\n\
+import torch, os, datetime, io\n\
+from diffusers import StableDiffusionPipeline\n\
+from flask_cors import CORS\n\
+\n\
+app = Flask(__name__)\n\
+CORS(app)\n\
+\n\
+print("Loading SD Model on CPU...")\n\
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)\n\
+pipe.to("cpu")\n\
+\n\
+@app.route("/api/generate", methods=["POST"])\n\
+def proxy():\n\
+    user_key = request.headers.get("x-api-key", "")\n\
+    # Simple check for your key\n\
+    if user_key != "sk-sd-user-1" and user_key != "sk-ess4l0ri37":\n\
+        return jsonify({"error": "Unauthorized"}), 401\n\
+    \n\
+    data = request.json\n\
+    prompt = data.get("prompt", "a sunset")\n\
+    image = pipe(prompt, num_inference_steps=15).images[0]\n\
+    \n\
+    img_io = io.BytesIO()\n\
+    image.save(img_io, "PNG")\n\
+    img_io.seek(0)\n\
+    return send_file(img_io, mimetype="image/png")\n\
+\n\
+if __name__ == "__main__":\n\
+    app.run(host="0.0.0.0", port=7860)\n' > $HOME/app.py
+
+# 6. Expose port 7860
+EXPOSE 7860
+
+# 7. Start the app
+CMD ["python3", "app.py"]