import argparse from PIL import Image import os from src.flux.xflux_pipeline import XFluxPipeline def create_argparser(): parser = argparse.ArgumentParser() parser.add_argument( "--prompt", type=str, required=True, help="The input text prompt" ) parser.add_argument( "--neg_prompt", type=str, default="", help="The input text negative prompt" ) parser.add_argument( "--img_prompt", type=str, default=None, help="Path to input image prompt" ) parser.add_argument( "--neg_img_prompt", type=str, default=None, help="Path to input negative image prompt" ) parser.add_argument( "--ip_scale", type=float, default=1.0, help="Strength of input image prompt" ) parser.add_argument( "--neg_ip_scale", type=float, default=1.0, help="Strength of negative input image prompt" ) parser.add_argument( "--local_path", type=str, default=None, help="Local path to the model checkpoint (Controlnet)" ) parser.add_argument( "--repo_id", type=str, default=None, help="A HuggingFace repo id to download model (Controlnet)" ) parser.add_argument( "--name", type=str, default=None, help="A filename to download from HuggingFace" ) parser.add_argument( "--ip_repo_id", type=str, default=None, help="A HuggingFace repo id to download model (IP-Adapter)" ) parser.add_argument( "--ip_name", type=str, default=None, help="A IP-Adapter filename to download from HuggingFace" ) parser.add_argument( "--ip_local_path", type=str, default=None, help="Local path to the model checkpoint (IP-Adapter)" ) parser.add_argument( "--lora_repo_id", type=str, default=None, help="A HuggingFace repo id to download model (LoRA)" ) parser.add_argument( "--lora_name", type=str, default=None, help="A LoRA filename to download from HuggingFace" ) parser.add_argument( "--lora_local_path", type=str, default=None, help="Local path to the model checkpoint (Controlnet)" ) parser.add_argument( "--device", type=str, default="cuda", help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)" ) parser.add_argument( "--offload", action='store_true', help="Offload model to CPU when not in use" ) parser.add_argument( "--use_ip", action='store_true', help="Load IP model" ) parser.add_argument( "--use_lora", action='store_true', help="Load Lora model" ) parser.add_argument( "--use_controlnet", action='store_true', help="Load Controlnet model" ) parser.add_argument( "--num_images_per_prompt", type=int, default=1, help="The number of images to generate per prompt" ) parser.add_argument( "--image", type=str, default=None, help="Path to image" ) parser.add_argument( "--lora_weight", type=float, default=0.9, help="Lora model strength (from 0 to 1.0)" ) parser.add_argument( "--control_weight", type=float, default=0.8, help="Controlnet model strength (from 0 to 1.0)" ) parser.add_argument( "--control_type", type=str, default="canny", choices=("canny", "openpose", "depth", "hed", "hough", "tile"), help="Name of controlnet condition, example: canny" ) parser.add_argument( "--model_type", type=str, default="flux-dev", choices=("flux-dev", "flux-dev-fp8", "flux-schnell"), help="Model type to use (flux-dev, flux-dev-fp8, flux-schnell)" ) parser.add_argument( "--width", type=int, default=1024, help="The width for generated image" ) parser.add_argument( "--height", type=int, default=1024, help="The height for generated image" ) parser.add_argument( "--num_steps", type=int, default=25, help="The num_steps for diffusion process" ) parser.add_argument( "--guidance", type=float, default=4, help="The guidance for diffusion process" ) parser.add_argument( "--seed", type=int, default=123456789, help="A seed for reproducible inference" ) parser.add_argument( "--true_gs", type=float, default=3.5, help="true guidance" ) parser.add_argument( "--timestep_to_start_cfg", type=int, default=5, help="timestep to start true guidance" ) parser.add_argument( "--save_path", type=str, default='results', help="Path to save" ) return parser def main(args): if args.image: image = Image.open(args.image) else: image = None xflux_pipeline = XFluxPipeline(args.model_type, args.device, args.offload) if args.use_ip: print('load ip-adapter:', args.ip_local_path, args.ip_repo_id, args.ip_name) xflux_pipeline.set_ip(args.ip_local_path, args.ip_repo_id, args.ip_name) if args.use_lora: print('load lora:', args.lora_local_path, args.lora_repo_id, args.lora_name) xflux_pipeline.set_lora(args.lora_local_path, args.lora_repo_id, args.lora_name, args.lora_weight) if args.use_controlnet: print('load controlnet:', args.local_path, args.repo_id, args.name) xflux_pipeline.set_controlnet(args.control_type, args.local_path, args.repo_id, args.name) image_prompt = Image.open(args.img_prompt) if args.img_prompt else None neg_image_prompt = Image.open(args.neg_img_prompt) if args.neg_img_prompt else None for _ in range(args.num_images_per_prompt): result = xflux_pipeline( prompt=args.prompt, controlnet_image=image, width=args.width, height=args.height, guidance=args.guidance, num_steps=args.num_steps, seed=args.seed, true_gs=args.true_gs, control_weight=args.control_weight, neg_prompt=args.neg_prompt, timestep_to_start_cfg=args.timestep_to_start_cfg, image_prompt=image_prompt, neg_image_prompt=neg_image_prompt, ip_scale=args.ip_scale, neg_ip_scale=args.neg_ip_scale, ) if not os.path.exists(args.save_path): os.mkdir(args.save_path) ind = len(os.listdir(args.save_path)) result.save(os.path.join(args.save_path, f"result_{ind}.png")) args.seed = args.seed + 1 if __name__ == "__main__": args = create_argparser().parse_args() main(args)