Update Dockerfile

Dockerfile

@@ -1,18 +1,27 @@
 FROM python:3.10-slim
 
+ENV DEBIAN_FRONTEND=noninteractive
+
 # ---------------- System deps ----------------
-RUN apt-get update && apt-get install -y \
+RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     libgl1 \
     libglib2.0-0 \
+    ca-certificates \
     && rm -rf /var/lib/apt/lists/*
 
 # ---------------- Python deps ----------------
-RUN pip install --no-cache-dir --upgrade \
-    torch \
-    torchvision \
-    torchaudio \
-    diffusers["torch"] \
+# Install torch FIRST, pinned, CPU by default (much smaller + stable)
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir \
+    torch==2.1.2 \
+    torchvision==0.16.2 \
+    torchaudio==2.1.2 \
+    --index-url https://download.pytorch.org/whl/cpu
+
+# Then the rest
+RUN pip install --no-cache-dir \
+    diffusers[torch] \
     transformers \
     accelerate \
     safetensors \
@@ -23,12 +32,14 @@ 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 
+ENV OMP_NUM_THREADS=1
+ENV MKL_NUM_THREADS=1
+ENV PYTORCH_ENABLE_MPS_FALLBACK=1
+
 # ---------------- Storage ----------------
 RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# ---------------- Python Application ----------------
+# ---------------- App ----------------
 RUN cat <<'EOF' > /app.py
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
@@ -39,7 +50,6 @@ 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"
@@ -48,31 +58,38 @@ LIMITS_PATH = f"{BASE}/limits.json"
 DEFAULT_LIMIT = 500
 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 --------
 print(f"Loading {MODEL_ID}...")
-pipe = DiffusionPipeline.from_pretrained(MODEL_ID)
+
+torch.set_grad_enabled(False)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+pipe = DiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    safety_checker=None
+)
+
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 
 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.")
+# ---- SPEED + MEMORY OPTS ----
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+
+if device == "cuda":
+    pipe.enable_model_cpu_offload()
+    pipe.unet.to(memory_format=torch.channels_last)
 
-print(f"Model loaded on {device}")
+print(f"Model ready on {device}")
 
-# -------- Helpers --------
 def whitelist():
     try:
         return set(open(WL_PATH).read().split())
@@ -88,7 +105,6 @@ def load_json(path):
 def save_json(path, data):
     json.dump(data, open(path, "w"))
 
-# -------- Routes --------
 @app.route("/", methods=["GET"])
 def health():
     return "LCM Image API Running", 200
@@ -99,7 +115,7 @@ def generate_key():
     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 data.get("unlimited") else int(data.get("limit", DEFAULT_LIMIT))
     save_json(LIMITS_PATH, limits)
@@ -113,29 +129,27 @@ def generate():
 
     data = request.get_json() or {}
     prompt = data.get("prompt", "").strip()
-    if not prompt: return jsonify({"error": "Prompt required"}), 400
+    if not prompt:
+        return jsonify({"error": "Prompt required"}), 400
 
-    # Rate Limiting
     limits = load_json(LIMITS_PATH)
     usage = load_json(USAGE_PATH)
     limit = limits.get(key, DEFAULT_LIMIT)
-    
+
     from datetime import datetime
     month = datetime.now().strftime("%Y-%m")
     used = usage.get(key, {}).get(month, 0)
-    
+
     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=4,
             guidance_scale=1.5
         ).images[0]
-        
+
         usage.setdefault(key, {})[month] = used + 1
         save_json(USAGE_PATH, usage)
 
@@ -143,8 +157,8 @@ def generate():
         image.save(buf, format="PNG")
         buf.seek(0)
         return send_file(buf, mimetype="image/png")
+
     except Exception as e:
-        print(f"Error: {e}")
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
@@ -152,12 +166,7 @@ if __name__ == "__main__":
 EOF
 
 # ---------------- Start ----------------
-RUN cat <<'EOF' > /start.sh
-#!/bin/bash
-python3 /app.py
-EOF
-
-RUN chmod +x /start.sh
+RUN echo '#!/bin/bash\npython3 /app.py' > /start.sh && chmod +x /start.sh
 
 EXPOSE 7860
-ENTRYPOINT ["/bin/bash", "/start.sh"]
+ENTRYPOINT ["/bin/bash", "/start.sh"]