initial commit

.dockerignore

@@ -0,0 +1,4 @@
+venv
+__pycache__
+.git
+outputs

.gitignore

@@ -0,0 +1,8 @@
+models/
+outputs/
+
+.venv/
+__pycache__/
+.env
+
+*.pyc

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app ./app
+
+RUN mkdir -p outputs
+RUN mkdir -p models
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+

README.md

@@ -0,0 +1,12 @@
+---
+title: Tiny SD FastAPI
+sdk: docker
+app_port: 7860
+emoji: 🎨
+colorFrom: blue
+colorTo: green
+---
+
+# Tiny SD FastAPI
+
+FastAPI + Docker + Tiny Stable Diffusion deployment.

app/main.py

@@ -0,0 +1,53 @@
+from fastapi import FastAPI, Query
+from fastapi.responses import FileResponse
+from pathlib import Path
+from datetime import datetime
+import re
+
+from app.model import generate_image
+
+app = FastAPI(
+    title="Tiny SD API",
+    docs_url="/docs",
+    redoc_url="/redoc",
+    openapi_url="/openapi.json",
+    swagger_ui_parameters={
+        "tryItOutEnabled": True
+    }
+)
+
+
+@app.get("/")
+def root():
+    return {
+        "message": "Tiny SD API is running"
+    }
+
+    
+@app.get("/generate")
+def generate(prompt: str):
+    image = generate_image(prompt)
+
+    output_dir = Path("outputs")
+    output_dir.mkdir(exist_ok=True)
+
+    safe_prompt = re.sub(
+        r'[\\/*?:"<>|]',
+        "_",
+        prompt
+    )[:40]
+
+    filename = (
+        f"{safe_prompt}-"
+        f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+    )
+
+    path = output_dir / filename
+
+    image.save(path)
+
+    return FileResponse(
+        path=path,
+        media_type="image/png",
+        filename=filename
+    )

app/model.py

@@ -0,0 +1,29 @@
+from diffusers import StableDiffusionPipeline
+import torch
+
+MODEL_NAME = "segmind/tiny-sd"
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    MODEL_NAME,
+    cache_dir="./models",
+    # local_files_only=True,
+    torch_dtype=torch.float32,
+    safety_checker=None
+)
+
+pipe.enable_attention_slicing()
+
+pipe.to("cpu")
+
+
+@torch.inference_mode()
+def generate_image(prompt):
+
+    image = pipe(
+        prompt,
+        height=256,
+        width=256,
+        num_inference_steps=10
+    ).images[0]
+
+    return image