Update Dockerfile

Dockerfile

@@ -1,58 +1,39 @@
-# 1. Use a more stable Python image
-FROM python:3.10-slim
+# Use an official PyTorch image with CUDA support
+FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
 
-# 2. System dependencies
-RUN apt-get update && apt-get install -y \
-    libgl1-mesa-glx libglib2.0-0 procps curl && \
-    rm -rf /var/lib/apt/lists/*
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    GRADIO_ALLOW_FLAGGING=never \
+    GRADIO_NUM_PORTS=1 \
+    GRADIO_SERVER_NAME="0.0.0.0" \
+    GRADIO_THEME=huggingface \
+    HOME=/home/user
 
-# 3. Python dependencies
-RUN pip3 install --no-cache-dir \
-    flask flask-cors requests diffusers transformers accelerate torch \
-    --extra-index-url https://download.pytorch.org/whl/cpu
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
 
-# 4. Set up non-root user (Required for HF)
+# Create a non-root user
 RUN useradd -m -u 1000 user
 USER user
-ENV HOME=/home/user
-ENV PATH=/home/user/.local/bin:$PATH
-WORKDIR $HOME
-
-# 5. Create the app file directly via a more robust method
-# Using printf instead of cat <<EOF to avoid shell/indentation errors
-RUN printf 'from flask import Flask, request, jsonify, send_file\n\
-import torch, os, datetime, io\n\
-from diffusers import StableDiffusionPipeline\n\
-from flask_cors import CORS\n\
-\n\
-app = Flask(__name__)\n\
-CORS(app)\n\
-\n\
-print("Loading SD Model on CPU...")\n\
-pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)\n\
-pipe.to("cpu")\n\
-\n\
-@app.route("/api/generate", methods=["POST"])\n\
-def proxy():\n\
-    user_key = request.headers.get("x-api-key", "")\n\
-    # Simple check for your key\n\
-    if user_key != "sk-sd-user-1" and user_key != "sk-ess4l0ri37":\n\
-        return jsonify({"error": "Unauthorized"}), 401\n\
-    \n\
-    data = request.json\n\
-    prompt = data.get("prompt", "a sunset")\n\
-    image = pipe(prompt, num_inference_steps=15).images[0]\n\
-    \n\
-    img_io = io.BytesIO()\n\
-    image.save(img_io, "PNG")\n\
-    img_io.seek(0)\n\
-    return send_file(img_io, mimetype="image/png")\n\
-\n\
-if __name__ == "__main__":\n\
-    app.run(host="0.0.0.0", port=7860)\n' > $HOME/app.py
-
-# 6. Expose port 7860
+WORKDIR $HOME/app
+
+# Copy requirements and install
+# (Assuming you have a requirements.txt with diffusers, transformers, accelerate, gradio)
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application
+COPY --chown=user . .
+
+# Make the startup script executable
+RUN chmod +x start.sh
+
+# Expose the port Gradio runs on
 EXPOSE 7860
 
-# 7. Start the app
-CMD ["python3", "app.py"]
+# Run the startup script
+CMD ["./start.sh"]