base

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.9
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8083
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8083"]
+
+

main.py

@@ -0,0 +1,46 @@
+from fastapi import FastAPI, Response
+from fastapi.middleware.cors import CORSMiddleware
+import torch
+from torch.cuda.amp import autocast
+from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline
+from io import BytesIO
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_credentials=True,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"]
+)
+
+
+
+# MODEL 1.4
+model_id = "CompVis/stable-diffusion-v1-4"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id, torch_dtype=torch.float32)
+
+pipe = pipe.to(device)
+
+# MODEL 1.0
+# model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, use_safetensors=True, variant="fp16")
+
+# pipe = pipe.to(device)
+
+@app.get("/")
+def generate(prompt: str):
+    with autocast(device):
+        image = pipe(prompt, guidance_scale=8.5).images[0]
+
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    buffer.seek(0)
+
+    return Response(content=buffer.getvalue(), media_type="image/png")

requirements.txt

@@ -0,0 +1,16 @@
+# fastapi
+# uvicorn
+# torch
+# transformers
+# diffusers
+
+fastapi
+torch
+diffusers
+transformers
+accelerate
+uvicorn
+datasets
+torch
+soundfile
+sentencepiece