[+] Main launcher

hi_diffusers/__main__.py

@@ -0,0 +1,40 @@
+from .nf4 import *
+
+import argparse
+import time
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument("prompt", type=str, help="Prompt to generate image from")
+    
+    parser.add_argument("-m", "--model", type=str, default="dev",
+                        help="Model to use",
+                        choices=["dev", "full", "fast"])
+    
+    parser.add_argument("-s", "--seed", type=int, default=-1, 
+                        help="Seed for generation")
+    
+    parser.add_argument("-r", "--res", type=str, default="1024x1024", 
+                        help="Resolution for generation", 
+                        choices=["1024x1024", "768x1360", "1360x768", "880x1168", "1168x880", "1248x832", "832x1248"])
+    
+    parser.add_argument("-o", "--output", type=str, default="output.png")
+    
+    args = parser.parse_args()
+    model_type = args.model
+    
+    # Initialize with default model
+    print(f"Loading model {model_type}...")
+    pipe, _ = load_models(model_type)
+    print("Model loaded successfully!")
+    
+    st = time.time()
+    
+    resolution = tuple(map(int, args.res.strip().split("x")))
+    image, seed = generate_image(pipe, model_type, args.prompt, resolution, args.seed)
+    image.save(args.output)
+    
+    print(f"Image saved to {args.output}, elapsed time: {time.time() - st:.2f} seconds")
+    print(f"Seed used: {seed}")

hi_diffusers/nf4.py

@@ -1,5 +1,4 @@
 import torch
-import argparse
 from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
 
 from . import HiDreamImagePipeline