From 1f178ff452c2f7f0d71813d7f7368ab819c7e8df Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Tue, 8 Apr 2025 22:11:36 +0000
Subject: [PATCH] [O] Better logging

---
 hi_diffusers/__main__.py |  5 ++++-
 hi_diffusers/nf4.py      | 11 +++++----
 hi_diffusers/shell.py    | 48 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 7 deletions(-)
 create mode 100644 hi_diffusers/shell.py

diff --git a/hi_diffusers/__main__.py b/hi_diffusers/__main__.py
index aa43440..9180bfa 100644
--- a/hi_diffusers/__main__.py
+++ b/hi_diffusers/__main__.py
@@ -1,10 +1,13 @@
 from .nf4 import *
 
 import argparse
-import time
+import Time
+import logging
 
 
 if __name__ == "__main__":
+    logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
+    
     parser = argparse.ArgumentParser()
     
     parser.add_argument("prompt", type=str, help="Prompt to generate image from")
diff --git a/hi_diffusers/nf4.py b/hi_diffusers/nf4.py
index e243593..2e05c99 100644
--- a/hi_diffusers/nf4.py
+++ b/hi_diffusers/nf4.py
@@ -38,15 +38,14 @@ MODEL_CONFIGS = {
 
 
 def log_vram(msg: str):
-    print(msg)
-    print(f"GPU memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+    print(f"{msg} (used {torch.cuda.memory_allocated() / 1024**2:.2f} MB VRAM)\n")
 
 
 def load_models(model_type: str):
     config = MODEL_CONFIGS[model_type]
     
     tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(LLAMA_MODEL_NAME)
-    log_vram("Tokenizer loaded!")
+    log_vram("✅ Tokenizer loaded!")
     
     text_encoder_4 = LlamaForCausalLM.from_pretrained(
         LLAMA_MODEL_NAME,
@@ -56,14 +55,14 @@ def load_models(model_type: str):
         torch_dtype=torch.bfloat16,
         device_map="auto",
     )
-    log_vram("Text encoder loaded!")
+    log_vram("✅ Text encoder loaded!")
 
     transformer = HiDreamImageTransformer2DModel.from_pretrained(
         config["path"],
         subfolder="transformer",
         torch_dtype=torch.bfloat16
     )
-    log_vram("Transformer loaded!")
+    log_vram("✅ Transformer loaded!")
     
     pipe = HiDreamImagePipeline.from_pretrained(
         config["path"],
@@ -73,7 +72,7 @@ def load_models(model_type: str):
         torch_dtype=torch.bfloat16,
     )
     pipe.transformer = transformer
-    log_vram("Pipeline loaded!")
+    log_vram("✅ Pipeline loaded!")
     pipe.enable_sequential_cpu_offload()
     
     return pipe, config
diff --git a/hi_diffusers/shell.py b/hi_diffusers/shell.py
new file mode 100644
index 0000000..bcfd12a
--- /dev/null
+++ b/hi_diffusers/shell.py
@@ -0,0 +1,48 @@
+from .nf4 import *
+
+import argparse
+import time
+import IPython
+import logging
+
+from IPython.display import Image, display
+
+
+def gen(prompt: str, seed: int = -1, res: str = "1024x1024", output="output.png"):
+    """Generate and display an image from the prompt."""
+    resolution = tuple(map(int, res.strip().split("x")))
+    
+    st = time.time()
+    image, final_seed = generate_image(pipe, args.model, prompt, resolution, seed)
+    image.save(output)
+    print(f"Image saved to {output}")
+    print(f"Seed used: {final_seed}, Time: {time.time() - st:.2f} seconds")
+    
+    # Display the image
+    display(Image(filename=output))
+
+
+if __name__ == "__main__":
+    logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-m", "--model", type=str, default="dev",
+                        help="Model to use",
+                        choices=["dev", "full", "fast"])
+    args = parser.parse_args()
+
+    # Load model 
+    print(f"Loading model {args.model}...")
+    pipe, _ = load_models(args.model)
+    print()
+    print("✅ Model loaded successfully!")
+    print("Try gen('your prompt here') to generate an image.")
+    print()
+    
+    # Set up IPython shell
+    banner = f"""
+HiDream-I1-nf4 Shell
+
+Model: {args.model} NF4 Quantized
+"""
+    IPython.start_ipython(argv=[], user_ns={"gen": gen}, banner=banner, display_banner=True)