From 89239d1c54b67dda883773fc65ba08615db8aa27 Mon Sep 17 00:00:00 2001
From: Eugene Brodsky <ebr@users.noreply.github.com>
Date: Sun, 12 Feb 2023 19:10:11 -0500
Subject: [PATCH] (updater) style 'pip' progress to use dark background

---
 ldm/invoke/config/invokeai_update.py | 54 +++++++++++++++++-----------
 1 file changed, 34 insertions(+), 20 deletions(-)

diff --git a/ldm/invoke/config/invokeai_update.py b/ldm/invoke/config/invokeai_update.py
index 8ad7290136..6e3bd6fbe1 100644
--- a/ldm/invoke/config/invokeai_update.py
+++ b/ldm/invoke/config/invokeai_update.py
@@ -2,16 +2,18 @@
 Minimalist updater script. Prompts user for the tag or branch to update to and runs
 pip install <path_to_git_source>.
 '''
-import os
+
 import platform
 import requests
+import subprocess
 from rich import box, print
-from rich.console import Console, Group, group
+from rich.console import Console, group
 from rich.panel import Panel
 from rich.prompt import Prompt
 from rich.style import Style
-from rich.syntax import Syntax
 from rich.text import Text
+from rich.live import Live
+from rich.table import Table
 
 from ldm.invoke import __version__
 
@@ -21,17 +23,19 @@ INVOKE_AI_REL="https://api.github.com/repos/invoke-ai/InvokeAI/releases"
 OS = platform.uname().system
 ARCH = platform.uname().machine
 
+ORANGE_ON_DARK_GREY = Style(bgcolor="grey23", color="orange1")
+
 if OS == "Windows":
     # Windows terminals look better without a background colour
     console = Console(style=Style(color="grey74"))
 else:
-    console = Console(style=Style(color="grey74", bgcolor="grey19"))
+    console = Console(style=Style(color="grey74", bgcolor="grey23"))
 
 def get_versions()->dict:
     return requests.get(url=INVOKE_AI_REL).json()
 
 def welcome(versions: dict):
-    
+
     @group()
     def text():
         yield f'InvokeAI Version: [bold yellow]{__version__}'
@@ -44,45 +48,55 @@ def welcome(versions: dict):
 [3] Manually enter the tag or branch name you wish to update'''
 
     console.rule()
-    print(
+    console.print(
         Panel(
             title="[bold wheat1]InvokeAI Updater",
             renderable=text(),
             box=box.DOUBLE,
             expand=True,
             padding=(1, 2),
-            style=Style(bgcolor="grey23", color="orange1"),
+            style=ORANGE_ON_DARK_GREY,
             subtitle=f"[bold grey39]{OS}-{ARCH}",
         )
     )
-    console.line()
+    # console.rule is used instead of console.line to maintain dark background
+    # on terminals where light background is the default
+    console.rule(characters=" ")
 
 def main():
     versions = get_versions()
     welcome(versions)
 
     tag = None
-    choice = Prompt.ask('Choice:',choices=['1','2','3'],default='1')
-    
+    choice = Prompt.ask(Text.from_markup(('[grey74 on grey23]Choice:')),choices=['1','2','3'],default='1')
+
     if choice=='1':
         tag = versions[0]['tag_name']
     elif choice=='2':
         tag = 'main'
     elif choice=='3':
-        tag = Prompt.ask('Enter an InvokeAI tag or branch name')
+        tag = Prompt.ask('[grey74 on grey23]Enter an InvokeAI tag or branch name')
+
+    console.print(Panel(f':crossed_fingers: Upgrading to [yellow]{tag}[/yellow]', box=box.MINIMAL, style=ORANGE_ON_DARK_GREY))
 
-    print(f':crossed_fingers: Upgrading to [yellow]{tag}[/yellow]')
     cmd = f'pip install {INVOKE_AI_SRC}/{tag}.zip --use-pep517'
-    print('')
-    print('')
-    if os.system(cmd)==0:
-        print(f':heavy_check_mark: Upgrade successful')
-    else:
-        print(f':exclamation: [bold red]Upgrade failed[/red bold]')
-    
+
+    progress = Table.grid(expand=True)
+    progress_panel = Panel(progress, box=box.MINIMAL, style=ORANGE_ON_DARK_GREY)
+
+    with subprocess.Popen(['bash', '-c', cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE) as proc:
+        progress.add_column()
+        with Live(progress_panel, console=console, vertical_overflow='visible'):
+            while proc.poll() is None:
+                for l in iter(proc.stdout.readline, b''):
+                    progress.add_row(l.decode().strip(), style=ORANGE_ON_DARK_GREY)
+        if proc.returncode == 0:
+            console.rule(f':heavy_check_mark: Upgrade successful')
+        else:
+            console.rule(f':exclamation: [bold red]Upgrade failed[/red bold]')
+
 if __name__ == "__main__":
     try:
         main()
     except KeyboardInterrupt:
         pass
-