cache

cache/local_cache.py

@@ -0,0 +1,34 @@
+from datetime import datetime, timedelta
+from functools import wraps
+
+CACHE_SIZE = 50
+
+_cache = {}
+_cache_time = {}
+
+
+def ttl_cache(func, key_name, ttl_secs=20):
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        # Assuming the prompt is the key for caching, change as necessary
+        key = kwargs.get(key_name)
+        ttl = timedelta(seconds=ttl_secs)
+        # Check cache
+        if key in _cache:
+            if datetime.now() - _cache_time[key] > ttl:
+                # Cache has expired
+                del _cache[key]
+                del _cache_time[key]
+            else:
+                return _cache[key]
+
+        # Call the actual function if not in cache or expired
+        response = await func(*args, **kwargs)
+
+        # Cache the result
+        _cache[key] = response
+        _cache_time[key] = datetime.now()
+
+        return response
+
+    return wrapper

routers/intference/stable_diffusion.py

@@ -6,6 +6,7 @@ from diffusers import DiffusionPipeline
 from fastapi import APIRouter
 from fastapi.responses import StreamingResponse
 
+from cache.local_cache import ttl_cache
 from config import settings
 
 router = APIRouter()
@@ -31,6 +32,7 @@ refiner.enable_attention_slicing()
 
 
 @router.get("/generate")
+@ttl_cache(key_name='prompt', ttl_secs=20)
 async def generate(prompt: str):
     """
     generate image
@@ -42,8 +44,6 @@ async def generate(prompt: str):
 
     prompt = f"single image. single model. {prompt}. zoomed in. full-body. real person. realistic. 4k. best quality."
     print(prompt)
-# "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly. bad anatomy, disfigured, poorly drawn face, mutation, mutated, extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, disconnected limbs, malformed hands, blurry, mutated hands and fingers, watermark, watermarked, oversaturated, censored, distorted hands, amputation, missing hands, obese, doubled face, double hands, two women, anime style, cartoon, toon. multiple people."
-    # prompt = "Designs should play with different textures and layering but stick to a monochrome palette. Think leather jackets over mesh tops, or satin draped over matte cotton. in a studio. zoomed-in."
 
     # run both experts
     image = base(