Upload main.py

main.py

@@ -804,7 +804,8 @@ def get_llama(query: str):
 @app.get("/ryuzaki/waifu")
 def waifu_pics(
     types: str="sfw",
-    category: str="neko"
+    category: str="neko",
+    media_type: str=None
 ):
     waifu_api = f"{SOURCE_WAIFU_URL}/{types}"
     waifu_param = f"{waifu_api}/{category}"
@@ -820,14 +821,11 @@ def waifu_pics(
         return f"Error request {e}"
     if waifu_image_url:
         try:
-            return {
-                "status": "true",
-                "randydev":{
-                    "image_url": waifu_image_url
-                }
-            }
-        except:
-            return {"status": "false", "message": "Error response"}
+            response_two = requests.get(waifu_image_url)
+            response_two.raise_for_status()
+        except requests.exceptions.RequestException:
+            raise HTTPException(status_code=500, detail="Internal server error")
+        return StreamingResponse(BytesIO(response_two.content), media_type=media_type)
     else:
         return {"status": "false", "message": "Error response."}