From ffbc2e788c113b278b9306aac4cea1d644eaf048 Mon Sep 17 00:00:00 2001 From: Zach Berwaldt Date: Wed, 30 Oct 2024 22:00:56 -0400 Subject: add logging, refactor generate image, keep track of prompts used. --- .gitignore | 2 +- prompts.txt | 2 ++ run_flux.py | 101 +++++++++++++++++++++++++++++++++++++++++++++++++++--------- 3 files changed, 89 insertions(+), 16 deletions(-) create mode 100644 prompts.txt diff --git a/.gitignore b/.gitignore index 2fa80d6..894a3f5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,2 @@ *.png - +*.log diff --git a/prompts.txt b/prompts.txt new file mode 100644 index 0000000..3268a51 --- /dev/null +++ b/prompts.txt @@ -0,0 +1,2 @@ +A space pilot, wearing her space flight suit, graying hair, determined to face the future. +Female with very large breasts. diff --git a/run_flux.py b/run_flux.py index 4a27816..89d188b 100644 --- a/run_flux.py +++ b/run_flux.py @@ -2,17 +2,42 @@ import torch import os import base64 import argparse +import argcomplete from diffusers import FluxPipeline, FluxImg2ImgPipeline from diffusers.utils import load_image from typing import Tuple from pathlib import Path from PIL import Image import uuid +import logging STORAGE_DIR: Path = Path.home() / "Pictures" STORAGE_DIR.mkdir(parents=True, exist_ok=True) +logger = logging.getLogger("run_flux") + +def image_completer(prefix, parsed_args, **kwargs): + image_dir = STORAGE_DIR / "Flux" + return [ + filename for filename in os.listdir(image_dir) + if filename.startswith(prefix) and os.path.isfile(os.path.join(image_dir, filename)) + ] + +def record_prompt(prompt, filename="prompts.txt"): + try: + with open(filename, "r") as file: + existing_prompts = set(line.strip() for line in file) + except FileNotFoundError: + existing_prompts = set() + + if prompt not in existing_prompts: + with open(filename, "a") as file: + file.write(prompt + "\n") + logger.info(f"Recording new prompt: \"{prompt}\"") + else: + logger.info(f"Prompt already exists in the file: \"{prompt}\"") + def load_flux_img_to_img(): pipeline = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipeline.enable_model_cpu_offload() @@ -27,17 +52,22 @@ def load_flux(): pipeline.vae.enable_tiling() return pipeline -def generate_image(pipeline, prompt, init_image=None, height=1024, width=1024, guideance_scale=0, num_images_per_prompt=1, num_inference_steps=50): - images = pipeline( - prompt=prompt, - guidance_scale=guideance_scale, - height=height, - width=width, - max_sequence_length=256, - num_inference_steps=num_inference_steps, - num_images_per_prompt=num_images_per_prompt, - image=init_image - ).images +def generate_image(pipeline, prompt, prompt_2=None, init_image=None, height=1024, width=1024, guideance_scale=0, num_images_per_prompt=1, num_inference_steps=50): + kwargs = { + "prompt": prompt, + "prompt_2": prompt_2, + "guidance_scale": guideance_scale, + "height":height, + "width": width, + "max_sequence_length": 256, + "num_inference_steps": num_inference_steps, + "num_images_per_prompt": num_images_per_prompt + } + + if isinstance(pipeline, FluxImg2ImgPipeline) and init_image is not None: + kwargs["image"] = init_image + + images = pipeline(**kwargs).images return images def generate_random_string(length=16) -> str: @@ -52,12 +82,25 @@ def parse_dimensions(dim_str: str) -> Tuple[int, int]: def main(): try: + logging.basicConfig(filename="flux.log", level=logging.INFO, format='%(asctime)s - %(levelname)s -> %(message)s', datefmt="%m/%d/%Y %I:%M:%S %p") + + logger.info("Parsing arguments") + args = parser.parse_args() - pipeline = load_flux_img_to_img() + logger.info("Choosing model...") + + pipeline = load_flux_img_to_img() if args.use_image else load_flux() + + if isinstance(pipeline, FluxPipeline): + logger.info("Using text-to-image model") + else: + logger.info("Using image-to-image model") + pipeline.to(torch.float16) - target_img = STORAGE_DIR / "1517481062292.jpg" + # target_img = STORAGE_DIR / "1517481062292.jpg" + target_img = STORAGE_DIR / "Flux" / "a23aae99-c8f1-4ce5-b91f-0b732774dadd.png" target_img_path = target_img.resolve(strict=True) @@ -67,25 +110,53 @@ def main(): width, height = args.size - images = generate_image(pipeline, init_image=init_image, prompt=args.prompt, width=width, height=height, guideance_scale=args.guideance_scale, num_images_per_prompt=args.number) + record_prompt(args.prompt) + + logger.info(f"Using prompt: \"{args.prompt}\"") + + logger.info("Generating image(s)...") + + images = generate_image( + pipeline=pipeline, + init_image=init_image, + prompt=args.prompt, + prompt_2=args.prompt2, + width=width, + height=height, + guideance_scale=args.guideance_scale, + num_images_per_prompt=args.number + ) + for image in images: filename = generate_random_string() filepath = STORAGE_DIR / "Flux" / f"{filename}.png" + logger.info(f"Saving {filepath}...") image.save(filepath) + + logger.info("Finished") except FileNotFoundError: print("\n Target image doesn't exist. Exiting...") exit(0) except KeyboardInterrupt: print('\nExiting early...') exit(0) + except Exception as e: + print(f"An error occured: {e}") + exit(1) parser = argparse.ArgumentParser(description="Generate some A.I. images", epilog="All done!") parser.add_argument("-n", "--number", type=int, default=1, help="the number of images you want to generate") parser.add_argument("-o", "--output", type=str, default="image", help="the name of the output image") parser.add_argument("-p", "--prompt", type=str, required=True, help="the prompt") +parser.add_argument("-p2", "--prompt2", type=str, help="A second prompt") parser.add_argument("-gs", "--guideance-scale", type=float, default=0) -parser.add_argument("--size", type=parse_dimensions, default="1024:1024") +parser.add_argument("--size", type=parse_dimensions, default="1024:1024", help="the size of the output images") +parser.add_argument("-u", "--use-image", action="store_true", help="use a predefined image") + +# parser.add_argument("-b", "--base-image").completer = image_completer + +# argcomplete.autocomplete(parser) if __name__ == "__main__": main() -- cgit v1.1