Enhanced Code

download.py

@@ -1,30 +1,40 @@
 import json
 from huggingface_hub import hf_hub_download
+from rich.console import Console
+from rich.table import Table
+from rich.prompt import Prompt, IntPrompt
+
+console = Console()
 
-# Read info.json
 with open("./lora/info.json", "r") as f:
     models = json.load(f)
 
-# Print menu
-print("Available models:")
+table = Table(title="Available Models", show_header=True, header_style="bold magenta")
+table.add_column("Index", style="dim", width=6)
+table.add_column("Name", style="bold cyan")
+table.add_column("Model", style="green")
+table.add_column("Description", style="yellow")
+
 for i, model in enumerate(models):
-    print(
-        f"[{i}] {model['name']} ({model['model'].split('/')[1]})\n> {model['description']}"
+    table.add_row(
+        str(i),
+        model['name'],
+        model['model'].split('/')[1],
+        model['description']
     )
 
-# Get user choice
-choice = int(input("\nSelect a model (enter number): "))
+console.print(table)
+
+choice = IntPrompt.ask("\nSelect a model (enter number)", choices=[str(i) for i in range(len(models))])
 
 if 0 <= choice < len(models):
     selected = models[choice]
     filename = f"{selected['name']}.safetensors"
 
-    print(f"\nDownloading {selected['name']}...")
-    print(
-        "Downloaded weights to: "
-        + hf_hub_download(
-            repo_id=selected["huggingface"], filename=filename, local_dir="./lora"
-        )
+    console.print(f"\n[bold green]Downloading {selected['name']}...[/bold green]")
+    downloaded_path = hf_hub_download(
+        repo_id=selected["huggingface"], filename=filename, local_dir="./lora"
     )
+    console.print(f"[bold blue]Downloaded weights to: {downloaded_path}[/bold blue]")
 else:
-    print("Invalid selection")
+    console.print("[bold red]Invalid selection[/bold red]")

generateImage.py

@@ -3,7 +3,6 @@
 
 
 def generate_image(prompt: str, lora: str, width: int, height: int):
-    # Load the model
     flux = Flux1(
         model_config=ModelConfig.FLUX1_DEV,
         quantize=8,
@@ -13,7 +12,6 @@ def generate_image(prompt: str, lora: str, width: int, height: int):
         lora_scales=[1.0],
     )
 
-    # Generate an image
     image = flux.generate_image(
         seed=7,
         prompt=prompt,

genmoji.py

@@ -6,112 +6,53 @@
 import argparse
 import json
 
-
 def get_unique_path(base_path):
-    directory = os.path.dirname(base_path)
-    filename = os.path.basename(base_path)
-    name, ext = os.path.splitext(filename)
-
-    # Remove any existing numbers from name
-    base_name = name.split("-")[0]
-
+    directory, filename = os.path.split(base_path)
+    base_name, ext = os.path.splitext(filename)
+    base_name = base_name.split("-")[0]
     counter = 1
     while True:
         new_path = os.path.join(directory, f"{base_name}-{counter:03d}{ext}")
         if not os.path.exists(new_path):
             return new_path
         counter += 1
 
+def main(user_prompt, direct, height, width, upscale_factor, output_path="output/genmoji.png", lora="flux-dev"):
+    try:
+        with open("./lora/info.json", "r") as f:
+            models = json.load(f)
+    except FileNotFoundError:
+        print("Error: info.json not found.")
+        sys.exit(1)
 
-def main(
-    user_prompt: str,
-    direct: bool,
-    height: int,
-    width: int,
-    upscale_factor: int,
-    output_path: str = "output/genmoji.png",
-    lora: str = "flux-dev",
-):
-    with open("./lora/info.json", "r") as f:
-        models = json.load(f)
-        metaprompt = "open-genmoji"
-        found = False
-        for model in models:
-            if model["name"] == lora:
-                found = True
-                metaprompt = model["metaprompt"]
-                break
+    metaprompt = next((model["metaprompt"] for model in models if model["name"] == lora), None)
+    if not metaprompt:
+        print(f"Error: LoRA {lora} does not exist. Run 'python download.py' to view and download available LoRAs.")
+        sys.exit(1)
 
-        if not found:
-            print(
-                f"Error: LoRA {lora} does not exist. Run 'python download.py' to view and download available LoRAs."
-            )
-            sys.exit(1)
+    lora_path = f"lora/{lora}.safetensors"
+    if not os.path.exists(lora_path):
+        print(f"Error: LoRA {lora} is not downloaded. Please run 'python download.py' to download it.")
+        sys.exit(1)
 
-        # Check if the lora file exists
-        lora_path = f"lora/{lora}.safetensors"
-        if not os.path.exists(lora_path):
-            print(
-                f"Error: LoRA {lora} is not downloaded. Please run 'python download.py' to download it."
-            )
-            sys.exit(1)
-        if not direct:
-            # Get the response from the prompt assistant
-            prompt_response = get_prompt_response(user_prompt, metaprompt)
-            print("Prompt Created: " + prompt_response)
-        elif direct:
-            prompt_response = user_prompt
-            print("Original prompt used: " + prompt_response)
+    prompt_response = user_prompt if direct else get_prompt_response(user_prompt, metaprompt)
+    print(f"Prompt Created: {prompt_response}" if not direct else f"Original prompt used: {prompt_response}")
 
-    # Generate the image using the response from the prompt assistant
     image = generate_image(prompt_response, lora, width, height)
-
-    output_width, output_height = image.size
-    new_img = image.resize((output_width * upscale_factor, output_height * upscale_factor), Image.LANCZOS)
+    new_img = image.resize((image.width * upscale_factor, image.height * upscale_factor), Image.LANCZOS)
 
     output_path = get_unique_path(output_path)
-    # Create output directory if it doesn't exist
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
     new_img.save(output_path)
 
-
 if __name__ == "__main__":
-    # When function is called it creates an object that looks like this:
-    # Namespace(user_prompt='a squirrel holding an iphone', assist=True, lora='flux-dev',
-    # width=160, height=160, upscale=5)
-    # the args object can then be passed into different functions
-    # and its values can be accessed, for example, like this:
-    # `args.height` --> returns `160` (as an int)
-
     parser = argparse.ArgumentParser()
     parser.add_argument("user_prompt", type=str, help="Your prompt")
-    parser.add_argument("-d", "--direct",
-                        action="store_true", help="Do not use prompt assistant")
-    parser.add_argument("-l", "--lora",
-                        nargs="?", default="flux-dev",
-                        type=str, help="The LoRA to use")
-    parser.add_argument("-iw", "--width",
-                        nargs="?", default=160,
-                        type=int, help="Image width")
-    parser.add_argument("-ih", "--height",
-                        nargs="?", default=160,
-                        type=int, help="Image height")
-    parser.add_argument("-u", "--upscale",
-                        nargs="?", default=5,
-                        type=int, help="Upscale factor")
+    parser.add_argument("-d", "--direct", action="store_true", help="Do not use prompt assistant")
+    parser.add_argument("-l", "--lora", default="flux-dev", type=str, help="The LoRA to use")
+    parser.add_argument("-iw", "--width", default=160, type=int, help="Image width")
+    parser.add_argument("-ih", "--height", default=160, type=int, help="Image height")
+    parser.add_argument("-u", "--upscale", default=5, type=int, help="Upscale factor")
     args = parser.parse_args()
-    user_prompt = args.user_prompt
-    lora = args.lora
-    direct = args.direct
-    height = args.height
-    width = args.width
-    upscale_factor = args.upscale
 
-    main(
-        user_prompt=user_prompt,
-        lora=lora,
-        direct=direct,
-        height=height,
-        width=width,
-        upscale_factor=upscale_factor
-    )
+    main(args.user_prompt, args.direct, args.height, args.width, args.upscale, lora=args.lora)

promptAssistant.py

@@ -2,33 +2,20 @@
 import json
 import os
 
-
 def get_prompt_response(user_prompt: str, metaprompt: str) -> str:
-    # The URL where the local server is running
     url = "http://localhost:1234/v1/chat/completions"
-
-    # The headers to indicate that we are sending JSON data
     headers = {"Content-Type": "application/json"}
 
-    # The JSON data payload
-    # Read the content from METAPROMPT.md
-    with open(
-        f"{os.path.abspath(os.path.dirname(__file__))}/metaprompt/{metaprompt}.md", "r"
-    ) as file:
+    with open(f"{os.path.abspath(os.path.dirname(__file__))}/metaprompt/{metaprompt}.md", "r") as file:
         prompt_content = file.read()
 
-    # Append the user prompt
-    full_prompt = prompt_content + f'\n\nUSER PROMPT: "{user_prompt}"'
-
-    # get the (pre-made) conversation history and append the current full prompt
+    full_prompt = f'{prompt_content}\n\nUSER PROMPT: "{user_prompt}"'
     json_path = f"{os.path.abspath(os.path.dirname(__file__))}/metaprompt/{metaprompt}.json"
 
     try:
         with open(json_path, "r") as json_file:
             conversation_history = json.load(json_file)
-            print("Using pre-existing conversation history")
     except FileNotFoundError:
-        print("No pre existing conversation history")
         conversation_history = {"messages": []}
 
     conversation_history["messages"].append({"role": "user", "content": full_prompt})
@@ -40,21 +27,9 @@ def get_prompt_response(user_prompt: str, metaprompt: str) -> str:
         "stream": False,
     }
 
-    # Making the POST request to the local server
     response = requests.post(url, headers=headers, data=json.dumps(data))
 
-    # Checking if the request was successful
     if response.status_code == 200:
-        # Returning the response content
-        return (
-            response.json()
-            .get("choices")[0]
-            .get("message")
-            .get("content")
-            .replace("```", "")
-            .replace("\n", "")
-        )
+        return response.json()["choices"][0]["message"]["content"].replace("```", "").replace("\n", "")
     else:
-        raise Exception(
-            f"Failed to get response: {response.status_code}, {response.text}"
-        )
+        raise Exception(f"Failed to get response: {response.status_code}, {response.text}")

resize.py

@@ -1,36 +1,14 @@
+from pathlib import Path
 from PIL import Image
-import os
-import sys
-
-
-def get_resized_filename(input_path):
-    """Generate output filename by adding '-resized' before extension"""
-    base, ext = os.path.splitext(input_path)
-    return f"{base}-resized{ext}"
-
-
-def resize_image(input_path, scale=5):
-    """Resize image by given scale factor using Lanczos resampling"""
-    try:
-        # Open and resize image
-        with Image.open(input_path) as img:
-            width, height = img.size
-            new_img = img.resize((width * scale, height * scale), Image.LANCZOS)
-
-            # Save resized image
-            output_path = get_resized_filename(input_path)
-            new_img.save(output_path)
-            print(f"Saved resized image to: {output_path}")
-
-    except Exception as e:
-        print(f"Error processing image: {e}")
-        sys.exit(1)
+from typing import Union
 
+def resize_image(input_path: Union[str, Path], scale: int = 5) -> None:
+    img = Image.open(input_path)
+    output_path = Path(input_path).stem + "-resized" + Path(input_path).suffix
+    img.resize((d * scale for d in img.size), Image.LANCZOS).save(output_path)
 
 if __name__ == "__main__":
+    import sys
     if len(sys.argv) != 2:
-        print("Usage: python resize_image.py <image_path>")
-        sys.exit(1)
-
-    input_path = sys.argv[1]
-    resize_image(input_path)
+        sys.exit("Usage: python resize.py <image_path>")
+    resize_image(sys.argv[1])