cache-fix

cache/local_cache.py

@@ -1,5 +1,8 @@
 from datetime import datetime, timedelta
 from functools import wraps
+from io import BytesIO
+
+from fastapi.responses import StreamingResponse
 
 CACHE_SIZE = 50
 
@@ -7,8 +10,7 @@ _cache = {}
 _cache_time = {}
 
 
-def ttl_cache(key_name, ttl_secs=20):
-
+def ttl_cache(key_name, media_type=None, ttl_secs=20):
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
@@ -22,7 +24,10 @@ def ttl_cache(key_name, ttl_secs=20):
                     del _cache[key]
                     del _cache_time[key]
                 else:
-                    return _cache[key]
+                    if media_type == 'image/png':
+                        return StreamingResponse(BytesIO(_cache[key]), media_type=media_type)
+                    else:
+                        return _cache[key]
 
             # Call the actual function if not in cache or expired
             response = await func(*args, **kwargs)

routers/intference/stable_diffusion.py

@@ -32,7 +32,7 @@ refiner.enable_attention_slicing()
 
 
 @router.get("/generate")
-@ttl_cache(key_name='prompt', ttl_secs=20)
+@ttl_cache(key_name='prompt', media_type="image/png", ttl_secs=20)
 async def generate(prompt: str):
     """
     generate image