From f8c5f0f74ce3f1d5e5a410fd86fd6828577f91ee Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Sun, 2 Jul 2023 13:32:23 +0900
Subject: [PATCH] fix: misc

---
 adetailer/traceback.py | 40 +++++++++++++++++++++++++---------------
 scripts/!adetailer.py  | 17 ++++++++++-------
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/adetailer/traceback.py b/adetailer/traceback.py
index 3ebc444..46dc85f 100644
--- a/adetailer/traceback.py
+++ b/adetailer/traceback.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import io
-from typing import Any
+import platform
+import sys
+from typing import Any, Callable
 
 from rich.console import Console, Group
 from rich.panel import Panel
@@ -11,7 +13,7 @@ from rich.traceback import Traceback
 from adetailer.__version__ import __version__
 
 
-def processing(*args):
+def processing(*args: Any) -> dict[str, Any]:
     try:
         from modules.processing import (
             StableDiffusionProcessingImg2Img,
@@ -28,7 +30,7 @@ def processing(*args):
             p = arg
             break
 
-    if not p:
+    if p is None:
         return {}
 
     info = {
@@ -47,7 +49,7 @@ def processing(*args):
     return info
 
 
-def sd_models():
+def sd_models() -> dict[str, str]:
     try:
         from modules import shared
 
@@ -62,16 +64,16 @@ def sd_models():
     }
 
 
-def ad_args(*args):
-    args = [
+def ad_args(*args: Any) -> dict[str, str]:
+    ad_args = [
         arg
         for arg in args
         if isinstance(arg, dict) and arg.get("ad_model", "None") != "None"
     ]
-    if not args:
+    if not ad_args:
         return {}
 
-    arg0 = args[0]
+    arg0 = ad_args[0]
     return {
         "version": __version__,
         "ad_model": arg0["ad_model"],
@@ -81,10 +83,7 @@ def ad_args(*args):
     }
 
 
-def sys_info():
-    import platform
-    import sys
-
+def sys_info() -> dict[str, Any]:
     try:
         import launch
 
@@ -104,7 +103,7 @@ def sys_info():
 
 def get_table(title: str, data: dict[str, Any]) -> Table:
     table = Table(title=title, highlight=True)
-    table.add_column(" ", style="dim")
+    table.add_column(" ", justify="right", style="dim")
     table.add_column("Value")
     for key, value in data.items():
         if not isinstance(value, str):
@@ -114,12 +113,23 @@ def get_table(title: str, data: dict[str, Any]) -> Table:
     return table
 
 
-def rich_traceback(func):
+def force_terminal_value():
+    try:
+        from modules.shared import cmd_opts
+
+        return True if hasattr(cmd_opts, "skip_torch_cuda_test") else None
+    except Exception:
+        return None
+
+
+def rich_traceback(func: Callable) -> Callable:
+    force_terminal = force_terminal_value()
+
     def wrapper(*args, **kwargs):
         string = io.StringIO()
         width = Console().width
         width = width - 4 if width > 4 else None
-        console = Console(file=string, force_terminal=True, width=width)
+        console = Console(file=string, force_terminal=force_terminal, width=width)
         try:
             return func(*args, **kwargs)
         except Exception as e:
diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py
index a9cfaa4..461b9cf 100644
--- a/scripts/!adetailer.py
+++ b/scripts/!adetailer.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import io
 import os
 import platform
 import re
@@ -16,7 +15,6 @@ from typing import Any
 import gradio as gr
 import torch
 from rich import print
-from rich.console import Console
 
 import modules
 from adetailer import (
@@ -94,6 +92,9 @@ class AfterDetailerScript(scripts.Script):
         self.cn_script = None
         self.cn_latest_network = None
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}(version={__version__})"
+
     def title(self):
         return AFTER_DETAILER
 
@@ -451,13 +452,15 @@ class AfterDetailerScript(scripts.Script):
         i2i.negative_prompt = negative_prompt
 
     @staticmethod
-    def compare_prompt(p, processed):
+    def compare_prompt(p, processed, n: int = 0):
         if p.prompt != processed.all_prompts[0]:
-            print(f"[-] ADetailer: applied ad_prompt: {processed.all_prompts[0]!r}")
+            print(
+                f"[-] ADetailer: applied {ordinal(n + 1)} ad_prompt: {processed.all_prompts[0]!r}"
+            )
 
         if p.negative_prompt != processed.all_negative_prompts[0]:
             print(
-                f"[-] ADetailer: applied ad_negative_prompt: {processed.all_negative_prompts[0]!r}"
+                f"[-] ADetailer: applied {ordinal(n + 1)} ad_negative_prompt: {processed.all_negative_prompts[0]!r}"
             )
 
     def is_need_call_process(self, p) -> bool:
@@ -547,10 +550,10 @@ class AfterDetailerScript(scripts.Script):
                     processed = process_images(p2)
                 except NansException as e:
                     msg = f"[-] ADetailer: 'NansException' occurred with {ordinal(n + 1)} settings.\n{e}"
-                    print(msg)
+                    print(msg, file=sys.stderr)
                     return False
 
-                self.compare_prompt(p2, processed)
+                self.compare_prompt(p2, processed, n=n)
                 p2 = copy(i2i)
                 p2.init_images = [processed.images[0]]