44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
from .nf4 import *
|
|
|
|
import argparse
|
|
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")
|
|
|
|
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}")
|