Merge branch 'main' of https://huggingface.co/spaces/Pacicap/api_stable_diffusion

Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.10
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# Set the working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Copy your app code
+COPY --chown=user . /app
+
+# Start the FastAPI app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,61 @@
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+from diffusers import DiffusionPipeline
+import torch
+import uuid
+import os
+from PIL import Image
+from fastapi.staticfiles import StaticFiles
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Accept from all for now
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+hf_model_ids = {
+    "model1": "Pacicap/FineTuned_claude_StableDiffussion_2_1",
+    "model2": "Pacicap/FineTuned_gpt4o_StableDiffussion_2_1"
+}
+
+loaded_models = {}
+
+class PromptInput(BaseModel):
+    prompt: str
+    model: str
+
+@app.post("/generate")
+def generate(data: PromptInput, request: Request):
+    model_key = data.model
+
+    if model_key not in hf_model_ids:
+        return {"error": "Invalid model selected"}
+
+    model_id = hf_model_ids[model_key]
+
+    if model_key not in loaded_models:
+        pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float32
+        ).to("cpu")  # CPU-safe for Spaces
+        loaded_models[model_key] = pipe
+    else:
+        pipe = loaded_models[model_key]
+
+    image = pipe(data.prompt).images[0]
+
+    os.makedirs("generated", exist_ok=True)
+    filename = f"{uuid.uuid4().hex}.png"
+    filepath = os.path.join("generated", filename)
+    image.save(filepath)
+
+    return {
+        "url": f"{request.base_url}generated/{filename}"
+    }
+
+app.mount("/generated", StaticFiles(directory="generated"), name="generated")