db for cache

.gitignore

@@ -1,2 +1,4 @@
 venv
-gradio_cached_examples
+gradio_cached_examples
+__pycache__/
+cache/

app.py

@@ -1,6 +1,3 @@
-import os
-import random
-import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
@@ -10,23 +7,95 @@ from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
 from fastapi import FastAPI
 import uvicorn
-from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import RedirectResponse
+from fastapi.responses import RedirectResponse, StreamingResponse
+import io
+import os
+from pathlib import Path
+from db import Database
+import uuid
+import logging
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+
+MAX_SEED = np.iinfo(np.int32).max
+USE_TORCH_COMPILE = os.environ.get("USE_TORCH_COMPILE", "0") == "1"
+SPACE_ID = os.environ.get('SPACE_ID', '')
+
+DB_PATH = Path("/data/cache") if SPACE_ID else Path("./cache")
+IMGS_PATH = DB_PATH / "imgs"
+DB_PATH.mkdir(exist_ok=True, parents=True)
+IMGS_PATH.mkdir(exist_ok=True, parents=True)
+
+database = Database(DB_PATH)
+
+with database() as db:
+    cursor = db.cursor()
+    cursor.execute("SELECT * FROM cache")
+    print(list(cursor.fetchall()))
+
+dtype = torch.bfloat16
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if torch.cuda.is_available():
+    prior_pipeline = StableCascadePriorPipeline.from_pretrained(
+        "stabilityai/stable-cascade-prior", torch_dtype=dtype
+    )  # .to(device)
+    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained(
+        "stabilityai/stable-cascade", torch_dtype=dtype
+    )  # .to(device)
+    prior_pipeline.to(device)
+    decoder_pipeline.to(device)
+
+    if USE_TORCH_COMPILE:
+        prior_pipeline.prior = torch.compile(
+            prior_pipeline.prior, mode="reduce-overhead", fullgraph=True
+        )
+        decoder_pipeline.decoder = torch.compile(
+            decoder_pipeline.decoder, mode="max-autotune", fullgraph=True
+        )
 
 
-class GenerateRequest(BaseModel):
-    prompt: str
-    negative_prompt: str = ""
-    seed: int = 0
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    prior_num_inference_steps: int = 20,
+    prior_guidance_scale: float = 4.0,
+    decoder_num_inference_steps: int = 10,
+    decoder_guidance_scale: float = 0.0,
+    num_images_per_prompt: int = 2,
+) -> PIL.Image.Image:
+
+    generator = torch.Generator().manual_seed(seed)
+    prior_output = prior_pipeline(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=prior_num_inference_steps,
+        timesteps=DEFAULT_STAGE_C_TIMESTEPS,
+        negative_prompt=negative_prompt,
+        guidance_scale=prior_guidance_scale,
+        num_images_per_prompt=num_images_per_prompt,
+        generator=generator,
+    )
+    decoder_output = decoder_pipeline(
+        image_embeddings=prior_output.image_embeddings,
+        prompt=prompt,
+        num_inference_steps=decoder_num_inference_steps,
+        # timesteps=decoder_timesteps,
+        guidance_scale=decoder_guidance_scale,
+        negative_prompt=negative_prompt,
+        generator=generator,
+        output_type="pil",
+    ).images
+
+    return decoder_output[0]
 
 
 app = FastAPI()
 origins = [
-    "http://localhost.tiangolo.com",
-    "https://localhost.tiangolo.com",
-    "http://localhost",
-    "http://localhost:8080",
+    "http://huggingface.co",
 ]
 
 app.add_middleware(
@@ -38,34 +107,37 @@ app.add_middleware(
 )
 
 
+@app.get("/image")
+async def generate_image(prompt: str, negative_prompt: str, seed: int = 2134213213):
+    cached_img = database.check(prompt, negative_prompt, seed)
+    if cached_img:
+        logging.info(f"Image found in cache: {cached_img[0]}")
+        return StreamingResponse(open(cached_img[0], "rb"), media_type="image/jpeg")
+
+    logging.info(f"Image not found in cache, generating new image")
+    pil_image = generate(prompt, negative_prompt, seed)
+    img_id = str(uuid.uuid4())
+    img_path = IMGS_PATH / f"{img_id}.jpg"
+    pil_image.save(img_path)
+    img_io = io.BytesIO()
+    pil_image.save(img_io, "JPEG")
+    img_io.seek(0)
+    database.insert(prompt, negative_prompt, str(img_path), seed)
+
+    return StreamingResponse(img_io, media_type="image/jpeg")
+
+
 @app.get("/")
 async def main():
     # redirect to https://huggingface.co/spaces/multimodalart/stable-cascade
-    return RedirectResponse("https://multimodalart-stable-cascade.hf.space/?__theme=system")
+    return RedirectResponse(
+        "https://multimodalart-stable-cascade.hf.space/?__theme=system"
+    )
 
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
 
-# MAX_SEED = np.iinfo(np.int32).max
-# USE_TORCH_COMPILE = False
-
-# dtype = torch.bfloat16
-# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-# if torch.cuda.is_available():
-#     prior_pipeline = StableCascadePriorPipeline.from_pretrained(
-#         "stabilityai/stable-cascade-prior", torch_dtype=dtype)  # .to(device)
-#     decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained(
-#         "stabilityai/stable-cascade",  torch_dtype=dtype)  # .to(device)
-#     prior_pipeline.to(device)
-#     decoder_pipeline.to(device)
-
-#     if USE_TORCH_COMPILE:
-#         prior_pipeline.prior = torch.compile(
-#             prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
-#         decoder_pipeline.decoder = torch.compile(
-#             decoder_pipeline.decoder, mode="max-autotune", fullgraph=True)
-
 
 # else:
 #     prior_pipeline = None

db.py

@@ -0,0 +1,54 @@
+import sqlite3
+from pathlib import Path
+
+
+class Database:
+    def __init__(self, db_path=None):
+        if db_path is None:
+            raise ValueError("db_path must be provided")
+        self.db_path = db_path
+        self.db_file = self.db_path / "cache.db"
+        if not self.db_file.exists():
+            print("Creating database")
+            print("DB_FILE", self.db_file)
+            db = sqlite3.connect(self.db_file)
+            with open(Path("schema.sql"), "r") as f:
+                db.executescript(f.read())
+            db.commit()
+            db.close()
+
+    def get_db(self):
+        db = sqlite3.connect(self.db_file, check_same_thread=False)
+        db.row_factory = sqlite3.Row
+        return db
+
+    def __enter__(self):
+        self.db = self.get_db()
+        return self.db
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.db.close()
+
+    def __call__(self):
+        return self
+
+    def insert(self, prompt: str, negative_prompt: str, image_path: str, seed: int):
+        with self() as db:
+            cursor = db.cursor()
+            cursor.execute(
+                "INSERT INTO cache (prompt, negative_prompt, image_path, seed) VALUES (?, ?, ?, ?)",
+                (prompt, negative_prompt, image_path, seed),
+            )
+            db.commit()
+
+    def check(self, prompt: str, negative_prompt: str, seed: int):
+        with self() as db:
+            cursor = db.cursor()
+            cursor.execute(
+                "SELECT image_path FROM cache WHERE prompt = ? AND negative_prompt = ?  AND seed = ? ORDER BY RANDOM() LIMIT 1",
+                (prompt, negative_prompt, seed),
+            )
+            image_path = cursor.fetchone()
+            if image_path:
+                return image_path
+            return False

requirements.txt

@@ -1,8 +1,9 @@
 git+https://github.com/kashif/diffusers.git@wuerstchen-v3
+fastapi==0.109.2
+numpy==1.26.4
+Pillow==10.2.0
+torch==2.2.0
+uvicorn==0.27.1
 accelerate
 safetensors
-transformers
-gradio
-fastapi
-pydantic
-uvicorn
+transformers

schema.sql

@@ -0,0 +1,14 @@
+PRAGMA foreign_keys = OFF;
+
+BEGIN TRANSACTION;
+
+CREATE TABLE cache (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    prompt TEXT NOT NULL,
+    negative_prompt TEXT NOT NULL,
+    image_path TEXT NOT NULL,
+    seed INTEGER NOT NULL,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+);
+
+COMMIT;