From 3333f94c3026b368e6bc0e3e2f9acb1d963e1d08 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sun, 26 Feb 2023 12:37:26 -0300
Subject: [PATCH] Make the gallery extension work on colab

---
 extensions/gallery/script.py | 38 +++++++++++++++++++++++++-----------
 modules/html_generator.py    |  2 +-
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/extensions/gallery/script.py b/extensions/gallery/script.py
index 8d606352..87216a15 100644
--- a/extensions/gallery/script.py
+++ b/extensions/gallery/script.py
@@ -4,6 +4,8 @@ from pathlib import Path
 
 import gradio as gr
 
+from modules.html_generator import image_to_base64
+
 
 def generate_html():
     css = """
@@ -29,10 +31,14 @@ def generate_html():
       background-color: gray;
     }
 
-    .character-gallery td {
-      text-align: center;
-      vertical-align: middle;
+    .character-gallery .image-td {
+      width: 150px;
     }
+
+    .character-gallery .character-td {
+      text-align: center !important;
+    }
+
     """
 
     table_html = f'<style>{css}</style><div class="character-gallery"><table>'
@@ -41,15 +47,25 @@ def generate_html():
     for file in Path("characters").glob("*"):
         if file.name.endswith(".json"):
             json_name = file.name
-            image_name = file.name.replace(".json", "")
+            character = file.name.replace(".json", "")
             table_html += "<tr>"
-            if Path(f"characters/{image_name}.png").exists():
-                image_html = f'<img src="file/characters/{image_name}.png">'
-            elif Path(f"characters/{image_name}.jpg").exists():
-                image_html = f'<img src="file/characters/{image_name}.jpg">'
-            else:
-                image_html = "<div class='placeholder'></div>"
-            table_html += f"<td>{image_html}</td><td>{image_name}</td>"
+            image_html = "<div class='placeholder'></div>"
+
+            for i in [
+                    f"characters/{character}.png",
+                    f"characters/{character}.jpg",
+                    f"characters/{character}.jpeg",
+                    ]:
+
+                path = Path(i)
+                if path.exists():
+                    try:
+                        image_html = f'<img src="data:image/png;base64,{image_to_base64(path)}">'
+                        break
+                    except:
+                        continue
+
+            table_html += f'<td class="image-td"=>{image_html}</td><td class="character-td">{character}</td>'
             table_html += "</tr>"
     table_html += "</table></div>"
 
diff --git a/modules/html_generator.py b/modules/html_generator.py
index 6e1fb8ac..dc2aa387 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -200,7 +200,7 @@ def image_to_base64(path):
 
     if (path in image_cache and mtime != image_cache[path][0]) or (path not in image_cache):
         img = Image.open(path)
-        img.thumbnail((100, 100))
+        img.thumbnail((200, 200))
         img_buffer = BytesIO()
         img.convert('RGB').save(img_buffer, format='PNG')
         image_cache[path] = [mtime, base64.b64encode(img_buffer.getvalue()).decode("utf-8")]