Upload main.py

main.py

@@ -805,7 +805,8 @@ def get_llama(query: str):
 def waifu_pics(
     types: str="sfw",
     category: str="neko",
-    media_type: str=None
+    media_type: str=None,
+    is_bytes: bool=False
 ):
     waifu_api = f"{SOURCE_WAIFU_URL}/{types}"
     waifu_param = f"{waifu_api}/{category}"
@@ -820,12 +821,20 @@ def waifu_pics(
     except Exception as e:
         return f"Error request {e}"
     if waifu_image_url:
-        try:
-            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)
+        if is_bytes:
+            try:
+                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": "true",
+                "randydev":{
+                    "image_url": waifu_image_url
+                }
+            }
     else:
         return {"status": "false", "message": "Error response."}