From 9229ed5e7e15608ec035d89cc5b2d19c3e629142 Mon Sep 17 00:00:00 2001
From: AlUlkesh <99896447+AlUlkesh@users.noreply.github.com>
Date: Fri, 14 Jul 2023 18:02:11 +0200
Subject: [PATCH] ImageReward fix, #214

---
 install.py               | 10 +++++++++-
 req_IR.txt               |  6 ++++++
 scripts/image_browser.py |  3 ++-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/install.py b/install.py
index cd16ab2..d87ff17 100644
--- a/install.py
+++ b/install.py
@@ -4,6 +4,14 @@ import os
 if not launch.is_installed("send2trash"):
     launch.run_pip("install Send2Trash", "Send2Trash requirement for image browser")
 
-if not launch.is_installed("ImageReward"):
+try:
+    import ImageReward
+    import dataset
+    import dill
+    import diffusers
+    import multiprocessing
+    import pyarrow
+    import xxhash
+except ImportError as e:
     req_IR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "req_IR.txt")
     launch.run_pip(f'install -r "{req_IR}" --no-deps image-reward', 'ImageReward requirement for image browser')
diff --git a/req_IR.txt b/req_IR.txt
index 81ef8ec..46b6466 100644
--- a/req_IR.txt
+++ b/req_IR.txt
@@ -1 +1,7 @@
+datasets
+diffusers
+dill
 fairscale
+multiprocess
+pyarrow
+xxhash
diff --git a/scripts/image_browser.py b/scripts/image_browser.py
index bb9ef67..1f681e2 100644
--- a/scripts/image_browser.py
+++ b/scripts/image_browser.py
@@ -48,8 +48,9 @@ except ImportError:
 try:
     import ImageReward
     image_reward_installed = True
-except ImportError:
+except ImportError as e:
     print("Image Browser: ImageReward is not installed, cannot be used.")
+    print(e)
     image_reward_installed = False
 
 # Force reload wib_db, as it doesn't get reloaded otherwise, if an extension update is started from webui