cache-fix

cache/local_cache.py

@@ -27,17 +27,14 @@ def ttl_cache(key_name, media_type=None, ttl_secs=20):
                     # if media_type == 'image/png':
                     #     return StreamingResponse(BytesIO(_cache[key]), media_type=media_type)
                     # else:
-                    return _cache[key]
+                    return StreamingResponse(BytesIO(_cache[key]), media_type="image/png")
 
             # Call the actual function if not in cache or expired
-            response: StreamingResponse = await func(*args, **kwargs)
+            response, image_data = await func(*args, **kwargs)
             # Cache the content of the response's body.
-            content = response.body
-            _cache[key] = content
+            _cache[key] = image_data
             _cache_time[key] = datetime.now()
 
-            # Reset the stream to its original state so it can be consumed by the caller
-            response.body.seek(0)
             return response
 
         return wrapper

routers/intference/stable_diffusion.py

@@ -63,5 +63,8 @@ async def generate(prompt: str):
 
     memory_stream = BytesIO()
     final_image.save(memory_stream, format="PNG")
+    image_data = memory_stream.getvalue()  # get bytes of the image
+
     memory_stream.seek(0)
-    return StreamingResponse(memory_stream, media_type="image/png")
+    return StreamingResponse(memory_stream, media_type="image/png"), image_data
+