Update Dockerfile

Dockerfile

@@ -1,34 +1,93 @@
-# Use an official PyTorch image with CUDA support
-FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
-
-# 1. Fix the "Time Zone" hang: Set to non-interactive and default to UTC
-ENV DEBIAN_FRONTEND=noninteractive \
-    TZ=Etc/UTC \
-    PYTHONUNBUFFERED=1 \
-    GRADIO_ALLOW_FLAGGING=never \
-    GRADIO_SERVER_NAME="0.0.0.0" \
-    HOME=/home/user
-
-# 2. Install system dependencies (tzdata is installed silently now)
+FROM python:3.10-slim
+
+# System deps (torch + diffusers need these or they sulk)
 RUN apt-get update && apt-get install -y \
-    tzdata \
     git \
-    libgl1-mesa-glx \
+    libgl1 \
     libglib2.0-0 \
     && rm -rf /var/lib/apt/lists/*
 
-# 3. Set up the Hugging Face User (Required for Spaces)
-RUN useradd -m -u 1000 user
-USER user
-WORKDIR $HOME/app
+# Python deps
+RUN pip install --no-cache-dir \
+    torch \
+    torchvision \
+    torchaudio \
+    diffusers \
+    transformers \
+    accelerate \
+    safetensors \
+    flask \
+    flask-cors \
+    pillow
+
+# Environment
+ENV HOME=/home/sd
+ENV HF_HOME=/home/sd/.cache
+ENV TRANSFORMERS_CACHE=/home/sd/.cache
+ENV DIFFUSERS_CACHE=/home/sd/.cache
+
+# Writable dirs (HF Spaces is picky)
+RUN mkdir -p /home/sd && chmod -R 777 /home/sd
+
+# -------- Flask Stable Diffusion 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
+from io import BytesIO
+import os
+
+app = Flask(__name__)
+CORS(app)
+
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+
+@app.route("/", methods=["GET"])
+def health():
+    return "Stable Diffusion API Running", 200
+
+@app.route("/api/generate", methods=["POST"])
+def generate():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    steps = int(data.get("steps", 25))
+    guidance = float(data.get("guidance", 7.5))
+
+    if not prompt:
+        return jsonify({"error": "No prompt provided"}), 400
+
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance
+    ).images[0]
+
+    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
+
+# -------- Startup Script --------
+RUN cat <<'EOF' > /start.sh
+#!/bin/bash
+echo "Starting Stable Diffusion API..."
+python3 /app.py
+EOF
 
-# 4. Install Python dependencies
-COPY --chown=user requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+RUN chmod +x /start.sh
 
-# 5. Copy the rest of your code
-COPY --chown=user . .
+EXPOSE 7860
 
-# 6. Make start script executable and launch
-RUN chmod +x start.sh
-CMD ["./start.sh"]
+ENTRYPOINT ["/bin/bash", "/start.sh"]