Commit fea1cca

@@ -1,8 +1,10 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+import base64
+import contextlib
+import json
 from pathlib import Path
 
-import anyio
 from pyrogram.client import Client
 from pyrogram.types import Message
 
@@ -43,11 +45,26 @@ async def cloudflare_text2img(client: Client, message: Message, model_id: str, p
             proxy=TEXT2IMG.CF_PROXY,
             rformat="content",
         )
-        if data := resp.get("content"):
-            path = Path(DOWNLOAD_DIR) / f"{rand_string(10)}.png"
-            async with await anyio.open_file(path, "wb") as f:
-                await f.write(data)
+        path = save_img(resp["content"])
+        if path.is_file():
             await send2tg(client, message, texts=f"{prompt}\n(By **{model_name}**)", media=[{"photo": path}], **kwargs)
             break
     await modify_progress(del_status=True, **kwargs)
     return {}
+
+
+def save_img(data: bytes) -> Path:
+    """Save image to file."""
+    path = Path(DOWNLOAD_DIR) / f"{rand_string(10)}.png"
+    # flux response
+    with contextlib.suppress(Exception):
+        json_data = json.loads(data.decode())
+        if json_data.get("success"):
+            with open(path, "wb") as f:
+                f.write(base64.b64decode(json_data["result"]["image"]))
+            return path
+
+    # stable diffusion response
+    with contextlib.suppress(Exception), open(path, "wb") as f:
+        f.write(data)
+    return path