fix error while returning result

app.py

@@ -83,7 +83,7 @@ def generate(
     upscale_by: float = 1.5,
     add_quality_tags: bool = True,
     progress=gr.Progress(track_tqdm=True),
-) -> Image:
+):
     generator = utils.seed_everything(seed)
 
     width, height = utils.aspect_ratio_handler(
@@ -132,6 +132,7 @@ def generate(
         }
     else:
         metadata["use_upscaler"] = None
+    metadata["model"] = "Animagine XL 3.0"
     logger.info(json.dumps(metadata, indent=4))
 
     try:
@@ -169,12 +170,15 @@ def generate(
                 output_type="pil",
             ).images
 
-        if images and IS_COLAB:
-            for image in images:
-                filepath = utils.save_image(image, metadata, OUTPUT_DIR)
-                logger.info(f"Image saved as {filepath} with metadata")
+        if images:
+            image_paths = [
+                utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB) for image in images
+            ]
 
-        return images, metadata
+            for image_path in image_paths:
+                logger.info(f"Image saved as {image_path} with metadata")
+
+        return image_paths, metadata
     except Exception as e:
         logger.exception(f"An error occurred: {e}")
         raise
@@ -221,14 +225,11 @@ with gr.Blocks(css="style.css") as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            run_button = gr.Button(
-                "Generate", 
-                variant="primary", 
-                scale=0
-            )
+            run_button = gr.Button("Generate", variant="primary", scale=0)
         result = gr.Gallery(
             label="Result", 
             columns=1, 
+            height="512px", 
             preview=True, 
             show_label=False
         )
@@ -239,10 +240,7 @@ with gr.Blocks(css="style.css") as demo:
             placeholder="Enter a negative prompt",
         )
         with gr.Row():
-            add_quality_tags = gr.Checkbox(
-                label="Add Quality Tags", 
-                value=True
-            )
+            add_quality_tags = gr.Checkbox(label="Add Quality Tags", value=True)
             quality_selector = gr.Dropdown(
                 label="Quality Tags Presets",
                 interactive=True,
@@ -348,49 +346,12 @@ with gr.Blocks(css="style.css") as demo:
         api_name=False,
     )
 
-    inputs = [
-        prompt,
-        negative_prompt,
-        seed,
-        custom_width,
-        custom_height,
-        guidance_scale,
-        num_inference_steps,
-        sampler,
-        aspect_ratio_selector,
-        style_selector,
-        quality_selector,
-        use_upscaler,
-        upscaler_strength,
-        upscale_by,
-        add_quality_tags,
-    ]
-
-    prompt.submit(
-        fn=utils.randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
-    ).then(
-        fn=generate,
-        inputs=inputs,
-        outputs=result,
-        api_name="run",
-    )
-    negative_prompt.submit(
-        fn=utils.randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
-    ).then(
-        fn=generate,
-        inputs=inputs,
-        outputs=result,
-        api_name=False,
-    )
-    run_button.click(
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
         fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
         outputs=seed,
@@ -398,8 +359,26 @@ with gr.Blocks(css="style.css") as demo:
         api_name=False,
     ).then(
         fn=generate,
-        inputs=inputs,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            custom_width,
+            custom_height,
+            guidance_scale,
+            num_inference_steps,
+            sampler,
+            aspect_ratio_selector,
+            style_selector,
+            quality_selector,
+            use_upscaler,
+            upscaler_strength,
+            upscale_by,
+            add_quality_tags,
+        ],
         outputs=[result, gr_metadata],
-        api_name=False,
+        api_name="run",
     )
-demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)

demo.ipynb

@@ -1,93 +1,114 @@
 {
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "538a3f0c-50c1-4952-9fcc-070d365c9a0f",
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import subprocess\n",
-    "from threading import Timer\n",
-    "from queue import Queue\n",
-    "\n",
-    "ROOT_DIR = \"/content\"\n",
-    "REPO_URL = \"https://huggingface.co/spaces/Linaqruf/animagine-xl\"\n",
-    "REPO_DIR = os.path.join(ROOT_DIR, \"cagliostro-webui\")\n",
-    "NGROK_TOKEN = \"\"\n",
-    "\n",
-    "os.environ[\"HF_TOKEN\"] = \"\"\n",
-    "os.environ[\"IS_COLAB\"] = \"1\"\n",
-    "os.environ[\"MODEL\"] = \"https://huggingface.co/cagliostrolab/animagine-xl-3.0/blob/main/animagine-xl-3.0.safetensors\"\n",
-    "os.environ[\"CACHE_EXAMPLES\"] = \"1\"\n",
-    "\n",
-    "def clone(url, dir, branch=None):\n",
-    "    subprocess.run([\"git\", \"clone\", url, dir], check=True)\n",
-    "    if branch:\n",
-    "        subprocess.run([\"git\", \"checkout\", branch], cwd=dir, check=True)\n",
-    "\n",
-    "def install_deps(dir):\n",
-    "    subprocess.run([\"pip\", \"install\", \"-r\", \"requirements.txt\"], cwd=dir, check=True)\n",
-    "\n",
-    "def ngrok_tunnel(port,queue,auth_token):\n",
-    "    ngrok.set_auth_token(auth_token)\n",
-    "    url = ngrok.connect(port)\n",
-    "    queue.put(url)\n",
-    "\n",
-    "def main():\n",
-    "    if not os.path.exists(REPO_DIR):\n",
-    "        print(f\"Cloning Repository to {REPO_DIR}\")\n",
-    "        clone(REPO_URL, REPO_DIR)\n",
-    "    print(f\"Installing required python libraries\")\n",
-    "    install_deps(REPO_DIR)\n",
-    "    print(\"Done!\")\n",
-    "\n",
-    "    os.chdir(REPO_DIR)\n",
-    "    \n",
-    "    if NGROK_TOKEN:\n",
-    "        try:\n",
-    "            from pyngrok import conf,ngrok\n",
-    "        except:\n",
-    "            !pip install -qqqq --upgrade setuptools\n",
-    "            !pip install -qqqq -U pyngrok\n",
-    "            from pyngrok import conf,ngrok\n",
-    "                                \n",
-    "        ngrok_output_queue = Queue()\n",
-    "        ngrok_thread = Timer(2, ngrok_tunnel, args=(7860, ngrok_output_queue, NGROK_TOKEN))\n",
-    "        ngrok_thread.start()\n",
-    "        ngrok_thread.join()\n",
-    "        \n",
-    "        print(ngrok_output_queue.get())  \n",
-    "    \n",
-    "    !python app.py\n",
-    "\n",
-    "if __name__ == \"__main__\":\n",
-    "    main()"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "538a3f0c-50c1-4952-9fcc-070d365c9a0f",
+      "metadata": {
+        "scrolled": true,
+        "id": "538a3f0c-50c1-4952-9fcc-070d365c9a0f"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "import subprocess\n",
+        "from threading import Timer\n",
+        "from queue import Queue\n",
+        "\n",
+        "def is_colab():\n",
+        "    try:\n",
+        "        import google.colab\n",
+        "        return True\n",
+        "    except ImportError:\n",
+        "        return False\n",
+        "\n",
+        "ROOT_DIR = \"/workspace/\" if not is_colab() else \"/content/\"\n",
+        "REPO_URL = \"https://huggingface.co/spaces/Linaqruf/animagine-xl\"\n",
+        "REPO_DIR = os.path.join(ROOT_DIR, \"animagine-xl\")\n",
+        "\n",
+        "NGROK_TOKEN = \"\"\n",
+        "NGROK_SUBDOMAIN = \"\"\n",
+        "PORT = 7860\n",
+        "\n",
+        "# os.environ[\"HF_TOKEN\"] = \"\"\n",
+        "os.environ[\"IS_COLAB\"] = \"1\"\n",
+        "os.environ[\"MODEL\"] = \"https://huggingface.co/cagliostrolab/animagine-xl-3.0/blob/main/animagine-xl-3.0.safetensors\"\n",
+        "os.environ[\"CACHE_EXAMPLES\"] = \"1\"\n",
+        "\n",
+        "def clone_repository(url, directory, branch=None):\n",
+        "    subprocess.run([\"git\", \"clone\", url, directory], check=True)\n",
+        "    if branch:\n",
+        "        subprocess.run([\"git\", \"checkout\", branch], cwd=directory, check=True)\n",
+        "\n",
+        "def install_dependencies(directory):\n",
+        "    dependencies = [\"accelerate==0.27.2\", \"diffusers==0.26.3\", \"gradio==4.20.0\",\n",
+        "                    \"invisible-watermark==0.2.0\", \"spaces==0.24.0\", \"omegaconf==2.3.0\", \"timm==0.9.10\"]\n",
+        "    if is_colab():\n",
+        "        subprocess.run([\"pip\", \"install\"] + dependencies, check=True)\n",
+        "    else:\n",
+        "        requirements_path = os.path.join(directory, \"requirements.txt\")\n",
+        "        subprocess.run([\"pip\", \"install\", \"-r\", requirements_path], check=True)\n",
+        "\n",
+        "def setup_ngrok_tunnel(port, queue, auth_token, subdomain):\n",
+        "    ngrok.set_auth_token(auth_token)\n",
+        "    url = ngrok.connect(port, bind_tls=True, subdomain=subdomain)\n",
+        "    queue.put(url)\n",
+        "\n",
+        "def main():\n",
+        "    if not os.path.exists(REPO_DIR):\n",
+        "        print(f\"Cloning repository to {REPO_DIR}\")\n",
+        "        clone_repository(REPO_URL, REPO_DIR)\n",
+        "\n",
+        "    print(\"Installing required Python libraries\")\n",
+        "    install_dependencies(REPO_DIR)\n",
+        "    print(\"Done!\")\n",
+        "\n",
+        "    os.chdir(REPO_DIR)\n",
+        "\n",
+        "    if NGROK_TOKEN:\n",
+        "        try:\n",
+        "            from pyngrok import conf, ngrok\n",
+        "        except ImportError:\n",
+        "            subprocess.run([\"pip\", \"install\", \"-qqqq\", \"--upgrade\", \"setuptools\"], check=True)\n",
+        "            subprocess.run([\"pip\", \"install\", \"-qqqq\", \"-U\", \"pyngrok\"], check=True)\n",
+        "            from pyngrok import conf, ngrok\n",
+        "\n",
+        "        ngrok.kill()\n",
+        "        ngrok_output_queue = Queue()\n",
+        "        ngrok_thread = Timer(2, setup_ngrok_tunnel, args=(PORT, ngrok_output_queue, NGROK_TOKEN, NGROK_SUBDOMAIN))\n",
+        "        ngrok_thread.start()\n",
+        "        ngrok_thread.join()\n",
+        "        print(ngrok_output_queue.get())\n",
+        "\n",
+        "    !python app.py\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "    main()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3 (ipykernel)",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.10.12"
+    },
+    "colab": {
+      "provenance": []
+    }
   },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.12"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

utils.py

@@ -4,6 +4,7 @@ import random
 import numpy as np
 import json
 import torch
+import uuid
 from PIL import Image, PngImagePlugin
 from datetime import datetime
 from dataclasses import dataclass
@@ -158,12 +159,14 @@ def preprocess_image_dimensions(width, height):
     return width, height
 
 
-def save_image(image, metadata, output_dir):
-    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+def save_image(image, metadata, output_dir, is_colab):
+    if is_colab:
+        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"image_{current_time}.png"   
+    else:
+        filename = str(uuid.uuid4()) + ".png"
     os.makedirs(output_dir, exist_ok=True)
-    filename = f"image_{current_time}.png"
     filepath = os.path.join(output_dir, filename)
-
     metadata_str = json.dumps(metadata)
     info = PngImagePlugin.PngInfo()
     info.add_text("metadata", metadata_str)