Update app.py

app.py

@@ -1,84 +1,60 @@
 import argparse
-import random
 from model import models
-from multit2i import (load_models, infer_fn, infer_rand_fn, save_gallery, change_model,
-                      warm_model, get_model_info_md, loaded_models, get_positive_prefix, get_positive_suffix,
-                      get_negative_prefix, get_negative_suffix, get_recom_prompt_type, set_recom_prompt_preset,
-                      randomize_seed, translate_to_en)
+from multit2i import (load_models, infer_fn, infer_rand_fn, save_gallery,
+    change_model, warm_model, get_model_info_md, loaded_models,
+    get_positive_prefix, get_positive_suffix, get_negative_prefix, get_negative_suffix,
+    get_recom_prompt_type, set_recom_prompt_preset, randomize_seed, translate_to_en)
 
-MAX_SEED = 2**32 - 1
-max_images = 6
+MAX_SEED = 2**32-1
+load_models(models)
 
-def generate_image(args):
-    load_models(models)
-    
-    model_name = args.model_name
+# Define a function to handle image generation
+def generate_images(args):
+    model_name = args.model
     prompt = args.prompt
     neg_prompt = args.neg_prompt
     width = args.width
     height = args.height
     steps = args.steps
     cfg = args.cfg
-    seed = args.seed if args.seed != -1 else random.randint(0, MAX_SEED)
-    
-    positive_prefix = args.positive_prefix
-    positive_suffix = args.positive_suffix
-    negative_prefix = args.negative_prefix
-    negative_suffix = args.negative_suffix
+    seed = args.seed
+    image_num = args.image_num
 
-    images = infer_fn(
-        model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
-        positive_prefix, positive_suffix, negative_prefix, negative_suffix
-    )
-    
-    save_gallery(images, args.output)
-
-def random_model_image(args):
-    load_models(models)
-    
-    model_name = random.choice(list(loaded_models.keys()))
-    prompt = args.prompt
-    neg_prompt = args.neg_prompt
-    width = args.width
-    height = args.height
-    steps = args.steps
-    cfg = args.cfg
-    seed = random.randint(0, MAX_SEED)
-    
     positive_prefix = args.positive_prefix
     positive_suffix = args.positive_suffix
     negative_prefix = args.negative_prefix
     negative_suffix = args.negative_suffix
 
-    images = infer_rand_fn(
-        model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
-        positive_prefix, positive_suffix, negative_prefix, negative_suffix
-    )
+    if seed == -1:
+        seed = randomize_seed()
+
+    # Infer function
+    for i in range(image_num):
+        output_image = infer_fn(model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
+                                positive_prefix, positive_suffix, negative_prefix, negative_suffix)
+        save_gallery(output_image, "gallery", f"output_image_{i}.png")
+
+# CLI Arguments parsing
+def main():
+    parser = argparse.ArgumentParser(description="Custom Image Generation CLI")
+    parser.add_argument("--model", type=str, required=True, help="Model name to use for generation.")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the image.")
+    parser.add_argument("--neg_prompt", type=str, default="", help="Negative prompt for the image.")
+    parser.add_argument("--width", type=int, default=0, help="Width of the image.")
+    parser.add_argument("--height", type=int, default=0, help="Height of the image.")
+    parser.add_argument("--steps", type=int, default=0, help="Number of inference steps.")
+    parser.add_argument("--cfg", type=float, default=0, help="Guidance scale.")
+    parser.add_argument("--seed", type=int, default=-1, help="Seed for randomization.")
+    parser.add_argument("--image_num", type=int, default=1, help="Number of images to generate.")
     
-    save_gallery(images, args.output)
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Image generation script using preloaded models.")
-
-    parser.add_argument('--model_name', type=str, default=list(loaded_models.keys())[0], help='The model to use for image generation.')
-    parser.add_argument('--prompt', type=str, required=True, help='The prompt for image generation.')
-    parser.add_argument('--neg_prompt', type=str, default="", help='Negative prompt to avoid unwanted elements.')
-    parser.add_argument('--width', type=int, default=0, help='Image width, 0 for default.')
-    parser.add_argument('--height', type=int, default=0, help='Image height, 0 for default.')
-    parser.add_argument('--steps', type=int, default=0, help='Number of inference steps.')
-    parser.add_argument('--cfg', type=float, default=0, help='Guidance scale (CFG).')
-    parser.add_argument('--seed', type=int, default=-1, help='Random seed, -1 for random.')
-    parser.add_argument('--output', type=str, default="output_gallery", help='Output directory for generated images.')
-    parser.add_argument('--positive_prefix', nargs='*', default=[], help='Positive prefix for the prompt.')
-    parser.add_argument('--positive_suffix', nargs='*', default=["Common"], help='Positive suffix for the prompt.')
-    parser.add_argument('--negative_prefix', nargs='*', default=[], help='Negative prefix for the prompt.')
-    parser.add_argument('--negative_suffix', nargs='*', default=["Common"], help='Negative suffix for the prompt.')
+    # Prefix and Suffix options
+    parser.add_argument("--positive_prefix", nargs="*", default=[], help="Positive prompt prefix.")
+    parser.add_argument("--positive_suffix", nargs="*", default=["Common"], help="Positive prompt suffix.")
+    parser.add_argument("--negative_prefix", nargs="*", default=[], help="Negative prompt prefix.")
+    parser.add_argument("--negative_suffix", nargs="*", default=["Common"], help="Negative prompt suffix.")
     
-    parser.add_argument('--random_model', action='store_true', help='Use random model for image generation.')
-
     args = parser.parse_args()
+    generate_images(args)
 
-    if args.random_model:
-        random_model_image(args)
-    else:
-        generate_image(args)
+if __name__ == "__main__":
+    main()