Update Dockerfile

Dockerfile

@@ -8,7 +8,6 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 # ---------------- Python deps ----------------
-# We upgrade these to ensure LCM support works
 RUN pip install --no-cache-dir --upgrade \
     torch \
     torchvision \
@@ -24,10 +23,12 @@ RUN pip install --no-cache-dir --upgrade \
 # ---------------- Env ----------------
 ENV HOME=/home/sd
 ENV HF_HOME=/home/sd/.cache
+# Limit threads to prevent CPU choking
+ENV OMP_NUM_THREADS=1 
 # ---------------- Storage ----------------
 RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# ---------------- Python Application (Written directly to file) ----------------
+# ---------------- Python Application ----------------
 RUN cat <<'EOF' > /app.py
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
@@ -45,9 +46,6 @@ 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 --------
@@ -56,16 +54,22 @@ 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 --------
+# -------- Load model --------
 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)
+
+# -------- CRITICAL MEMORY FIXES --------
+if device == "cpu":
+    # Slices attention computation into chunks (Saves ~2GB RAM)
+    pipe.enable_attention_slicing()
+    # Slices VAE decoding (Saves ~1GB RAM)
+    pipe.enable_vae_slicing()
+    print("Memory optimizations enabled for CPU.")
+
 print(f"Model loaded on {device}")
 
 # -------- Helpers --------
@@ -84,30 +88,23 @@ def load_json(path):
 def save_json(path, data):
     json.dump(data, open(path, "w"))
 
-# -------- Health --------
+# -------- Routes --------
 @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)
-
     with open(WL_PATH, "a") as f:
         f.write(key + "\n")
-
+    
     limits = load_json(LIMITS_PATH)
-    limits[key] = "unlimited" if unlimited else int(limit)
+    limits[key] = "unlimited" if data.get("unlimited") else int(data.get("limit", DEFAULT_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", "")
@@ -116,49 +113,40 @@ def generate():
 
     data = request.get_json() or {}
     prompt = data.get("prompt", "").strip()
+    if not prompt: return jsonify({"error": "Prompt required"}), 400
 
-    # 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
+    # Rate Limiting
     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:
+    if limit != "unlimited" and used >= limit:
         return jsonify({"error": "Monthly limit reached"}), 429
 
     # Generate
     try:
+        # Hardcoded for stability
         image = pipe(
             prompt=prompt,
-            num_inference_steps=steps,
-            guidance_scale=guidance
+            num_inference_steps=4,
+            guidance_scale=1.5
         ).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")
     except Exception as e:
-        print(f"Generation Error: {e}")
+        print(f"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