Update Dockerfile

Dockerfile

@@ -1,158 +1,48 @@
 FROM python:3.10-slim
 
-# ---------------- System deps ----------------
+# System dependencies
 RUN apt-get update && apt-get install -y \
     git \
     libgl1 \
     libglib2.0-0 \
+    wget \
     && rm -rf /var/lib/apt/lists/*
 
-# ---------------- Python deps ----------------
+# Optimized Python dependencies
 RUN pip install --no-cache-dir \
-    torch \
-    torchvision \
-    torchaudio \
-    diffusers["torch"] \
-    transformers \
-    accelerate \
-    safetensors \
-    flask \
-    flask-cors \
-    pillow
-
-# ---------------- Env ----------------
+    torch==2.1.2 \
+    torchvision==0.16.2 \
+    torchaudio==2.1.2 \
+    --index-url https://download.pytorch.org/whl/cu118 \
+    && pip install --no-cache-dir \
+    diffusers==0.26.3 \
+    transformers==4.38.2 \
+    accelerate==0.27.2 \
+    safetensors==0.4.2 \
+    flask==3.0.3 \
+    flask-cors==4.0.0 \
+    pillow==10.2.0 \
+    xformers==0.0.24
+
+# Environment variables for caching
 ENV HOME=/home/sd
-ENV HF_HOME=/home/sd/.cache
-ENV TRANSFORMERS_CACHE=/home/sd/.cache
-ENV DIFFUSERS_CACHE=/home/sd/.cache
+ENV HF_HOME=/home/sd/.cache/huggingface
+ENV TRANSFORMERS_CACHE=/home/sd/.cache/huggingface/models
+ENV DIFFUSERS_CACHE=/home/sd/.cache/huggingface/diffusers
+ENV PYTHONUNBUFFERED=1
+ENV HF_ENDPOINT=https://hf-mirror.com
 
-# ---------------- Storage ----------------
+# Create directory with proper permissions
 RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# ---------------- Flask API ----------------
-RUN cat <<'EOF' > /app.py
-from flask import Flask, request, jsonify, send_file
-from flask_cors import CORS
-from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
-import torch, os, json, secrets
-from io import BytesIO
-
-app = Flask(__name__)
-CORS(app)
-
-# -------- 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 = "stabilityai/stable-diffusion-xl-base-1.0"
-
-# -------- 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 = StableDiffusionXLPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
-# -------- Helpers --------
-def whitelist():
-    return set(open(WL_PATH).read().split())
-
-def load_json(path):
-    try:
-        return json.load(open(path))
-    except:
-        return {}
-
-def save_json(path, data):
-    json.dump(data, open(path, "w"))
-
-# -------- Health --------
-@app.route("/", methods=["GET"])
-def health():
-    return "Image XL API Running", 200
-
-# -------- Key generator --------
-@app.route("/generate-key", methods=["POST"])
-def generate_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_json(LIMITS_PATH)
-    limits[key] = "unlimited" if unlimited else int(limit)
-    save_json(LIMITS_PATH, limits)
-
-    return jsonify({"key": key, "limit": limits[key]})
+# Copy application code
+COPY app.py /app.py
 
-# -------- Image generation --------
-@app.route("/api/generate", methods=["POST"])
-def generate():
-    key = request.headers.get("x-api-key", "")
-    if key not in whitelist():
-        return jsonify({"error": "Unauthorized"}), 401
-
-    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")
-
-    from datetime import datetime
-    month = 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
-
-    image = pipe(
-        prompt=prompt,
-        num_inference_steps=steps,
-        guidance_scale=guidance
-    ).images[0]
-
-    usage.setdefault(key, {})[month] = used + 1
-    save_json(USAGE_PATH, usage)
-
-    buf = BytesIO()
-    image.save(buf, format="PNG")
-    buf.seek(0)
-
-    return send_file(buf, mimetype="image/png")
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
-EOF
-
-# ---------------- Start ----------------
-RUN cat <<'EOF' > /start.sh
-#!/bin/bash
-python3 /app.py
-EOF
-
-RUN chmod +x /start.sh
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
+    CMD python -c "import requests; requests.get('http://localhost:7860/', timeout=2)"
 
 EXPOSE 7860
-ENTRYPOINT ["/bin/bash", "/start.sh"]
+
+# Start command
+CMD ["python", "/app.py"]