Update Dockerfile

Dockerfile

@@ -1,48 +1,175 @@
 FROM python:3.10-slim
 
-# System dependencies
+# ---------------- System deps ----------------
 RUN apt-get update && apt-get install -y \
     git \
     libgl1 \
     libglib2.0-0 \
-    wget \
     && rm -rf /var/lib/apt/lists/*
 
-# Optimized Python dependencies
-RUN pip install --no-cache-dir \
-    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/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
+# ---------------- Python deps ----------------
+# We upgrade these to ensure LCM support works
+RUN pip install --no-cache-dir --upgrade \
+    torch \
+    torchvision \
+    torchaudio \
+    diffusers["torch"] \
+    transformers \
+    accelerate \
+    safetensors \
+    flask \
+    flask-cors \
+    pillow
 
-# Create directory with proper permissions
+# ---------------- Env ----------------
+ENV HOME=/home/sd
+ENV HF_HOME=/home/sd/.cache
+# ---------------- Storage ----------------
 RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# Copy application code
-COPY app.py /app.py
+# ---------------- Python Application (Written directly to file) ----------------
+RUN cat <<'EOF' > /app.py
+from flask import Flask, request, jsonify, send_file
+from flask_cors import CORS
+from diffusers import DiffusionPipeline, LCMScheduler
+import torch, os, json, secrets
+from io import BytesIO
 
-# Health check
-HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
-    CMD python -c "import requests; requests.get('http://localhost:7860/', timeout=2)"
+app = Flask(__name__)
+CORS(app)
 
-EXPOSE 7860
+# -------- 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 Config --------
+# LCM Dreamshaper is SD1.5 based (small) and needs only 4-8 steps (fast)
+MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7"
+
+# -------- 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 --------
+print(f"Loading {MODEL_ID}...")
+pipe = DiffusionPipeline.from_pretrained(MODEL_ID)
+
+# Ensure we use the LCM Scheduler (Fixes the IndexError crash)
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+
+# Detect hardware (CPU vs CUDA)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+print(f"Model loaded on {device}")
+
+# -------- Helpers --------
+def whitelist():
+    try:
+        return set(open(WL_PATH).read().split())
+    except:
+        return set()
+
+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 "LCM Image 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)
 
-# Start command
-CMD ["python", "/app.py"]
+    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]})
+
+# -------- 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()
+
+    # LCM Optimization:
+    # LCM works best between 4 and 8 steps. 
+    # We ignore the user's requested 'steps' to ensure speed and stability.
+    steps = 4 
+    guidance = 2.0 
+
+    if not prompt:
+        return jsonify({"error": "Prompt required"}), 400
+
+    # Rate Limiting Logic
+    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
+
+    # Generate
+    try:
+        image = pipe(
+            prompt=prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance
+        ).images[0]
+    except Exception as e:
+        print(f"Generation Error: {e}")
+        return jsonify({"error": str(e)}), 500
+
+    # Save Usage
+    usage.setdefault(key, {})[month] = used + 1
+    save_json(USAGE_PATH, usage)
+
+    # Return Image
+    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
+
+EXPOSE 7860
+ENTRYPOINT ["/bin/bash", "/start.sh"]