Upload 3 files

Dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.9
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+
+RUN useradd user
+
+USER user
+
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user:user . $HOME/app
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,36 @@
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from diffusers import StableDiffusionXLPipeline
+import torch
+import base64
+from io import BytesIO
+
+#create fastapi instance
+app = FastAPI()
+
+# Load the pre-trained Stable Diffusion XL model
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "segmind/SSD-1B",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+)
+pipe.to("cuda")
+# pipe.enable_xformers_memory_efficient_attention()  # for torch < 2.0
+
+
+@app.get("/")
+def home():
+    return {"message": "Stable Diffusion API running"}
+
+#Define function to generate image from text prompt
+@app.get("/generate-image/")
+def generate_image(prompt: str, negative_prompt: str = None):
+    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
+    
+    # Convert to base64 string
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return JSONResponse(content={"image_base64": img_str})

requirements.txt

@@ -0,0 +1,7 @@
+fastapi[standard]
+uvicorn[standard]
+torch
+transformers
+accelerate
+safetensors
+diffusers