Update Dockerfile

Dockerfile

@@ -12,7 +12,7 @@ RUN pip install --no-cache-dir \
     torch \
     torchvision \
     torchaudio \
-    diffusers \
+    diffusers["torch"] \
     transformers \
     accelerate \
     safetensors \
@@ -29,12 +29,12 @@ ENV DIFFUSERS_CACHE=/home/sd/.cache
 # ---------------- Storage ----------------
 RUN mkdir -p /home/sd && chmod -R 777 /home/sd
 
-# ---------------- Image Guard API ----------------
+# ---------------- Flask API ----------------
 RUN cat <<'EOF' > /app.py
 from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-from diffusers import StableDiffusionPipeline
-import torch, os, json, datetime, secrets
+from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
+import torch, os, json, secrets
 from io import BytesIO
 
 app = Flask(__name__)
@@ -47,7 +47,7 @@ USAGE_PATH = f"{BASE}/usage.json"
 LIMITS_PATH = f"{BASE}/limits.json"
 
 DEFAULT_LIMIT = 500
-MODEL_ID = "stable-diffusion-v1-5"
+MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
 
 # -------- Init storage --------
 os.makedirs(BASE, exist_ok=True)
@@ -56,10 +56,11 @@ for p in [WL_PATH, USAGE_PATH, LIMITS_PATH]:
         open(p, "w").write("{}" if p.endswith(".json") else "")
 
 # -------- Load model once --------
-pipe = StableDiffusionPipeline.from_pretrained(
+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 --------
@@ -78,7 +79,7 @@ def save_json(path, data):
 # -------- Health --------
 @app.route("/", methods=["GET"])
 def health():
-    return "Image API Running", 200
+    return "Image XL API Running", 200
 
 # -------- Key generator --------
 @app.route("/generate-key", methods=["POST"])
@@ -96,10 +97,7 @@ def generate_key():
     limits[key] = "unlimited" if unlimited else int(limit)
     save_json(LIMITS_PATH, limits)
 
-    return jsonify({
-        "key": key,
-        "limit": limits[key]
-    })
+    return jsonify({"key": key, "limit": limits[key]})
 
 # -------- Image generation --------
 @app.route("/api/generate", methods=["POST"])
@@ -122,7 +120,8 @@ def generate():
     limit = limits.get(key, DEFAULT_LIMIT)
     unlimited = (limit == "unlimited")
 
-    month = datetime.datetime.now().strftime("%Y-%m")
+    from datetime import datetime
+    month = datetime.now().strftime("%Y-%m")
     used = usage.get(key, {}).get(month, 0)
 
     if not unlimited and used >= limit: