fix error when handling outputs

app.py

@@ -11,10 +11,7 @@ import logging
 from PIL import Image, PngImagePlugin
 from datetime import datetime
 from diffusers.models import AutoencoderKL
-from diffusers import (
-    StableDiffusionXLPipeline, 
-    StableDiffusionXLImg2ImgPipeline
-)
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -31,7 +28,10 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
 
-MODEL = os.getenv("MODEL", "https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors")
+MODEL = os.getenv(
+    "MODEL",
+    "https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors",
+)
 
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
@@ -82,7 +82,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(
@@ -131,6 +131,7 @@ def generate(
         }
     else:
         metadata["use_upscaler"] = None
+    metadata["model"] = DESCRIPTION
     logger.info(json.dumps(metadata, indent=4))
 
     try:
@@ -167,13 +168,17 @@ def generate(
                 generator=generator,
                 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")
-            
-        return images, metadata
+
+        if images:
+            image_paths = [
+                utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
+                for image in images
+            ]
+
+            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
@@ -251,6 +256,22 @@ with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
                         value="896 x 1152",
                         container=True,
                     )
+                with gr.Group(visible=False) as custom_resolution:
+                    with gr.Row():
+                        custom_width = gr.Slider(
+                            label="Width",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                        custom_height = gr.Slider(
+                            label="Height",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
                 with gr.Group():
                     use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
                     with gr.Row() as upscaler_row:
@@ -270,22 +291,6 @@ with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
                             value=1.5,
                             visible=False,
                         )
-                with gr.Group(visible=False) as custom_resolution:
-                    with gr.Row():
-                        custom_width = gr.Slider(
-                            label="Width",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                        custom_height = gr.Slider(
-                            label="Height",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
                 with gr.Group():
                     sampler = gr.Dropdown(
                         label="Sampler",
@@ -320,6 +325,7 @@ with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
             result = gr.Gallery(
                 label="Result", 
                 columns=1, 
+                height='100%', 
                 preview=True, 
                 show_label=False
             )
@@ -347,25 +353,12 @@ with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") 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(
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
         fn=utils.randomize_seed_fn,
         inputs=[seed, randomize_seed],
         outputs=seed,
@@ -373,32 +366,26 @@ with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
         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(
-        fn=utils.randomize_seed_fn,
-        inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        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,110 +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",
-    "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/cagliostrolab/animagine-xl-3.1\"\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.1/blob/main/animagine-xl-3.1.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",
-    "    subprocess.run([\"python\", \"app.py\"], check=True)\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/cagliostrolab/animagine-xl-3.1\"\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.1/blob/main/animagine-xl-3.1.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)