import argparse import json import logging import os import sys from types import SimpleNamespace import gc import json import os import random import sys import time from functools import partial from glob import glob from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision.transforms import functional as TF from diffusers.models import AutoencoderKL from transformers import BertModel, BertTokenizer, logging as tf_logging # from hydit.constants import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER from hydit.lr_scheduler import WarmupLR from hydit.data_loader.arrow_load_stream import TextImageArrowStream from hydit.diffusion import create_diffusion from hydit.ds_config import deepspeed_config_from_args from hydit.modules.ema import EMA from hydit.modules.fp16_layers import Float16Module from hydit.modules.models import HUNYUAN_DIT_MODELS from hydit.modules.posemb_layers import init_image_posemb from hydit.utils.tools import set_seeds, get_trainable_params from IndexKits.index_kits import ResolutionGroup from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler from peft import LoraConfig, get_peft_model from tqdm import tqdm UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" def set_unet_path(path): global UNET_PATH UNET_PATH = path VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" def set_vae_ema_path(path): global VAE_EMA_PATH VAE_EMA_PATH = path TOKENIZER = "ckpts/t2i/tokenizer" def set_tokenizer_path(path): global TOKENIZER TOKENIZER = path TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' def set_text_encoder_path(path): global TEXT_ENCODER TEXT_ENCODER = path T5_ENCODER = { 'MT5': 'ckpts/t2i/mt5', 'attention_mask': True, 'layer_index': -1, 'attention_pool': True, 'torch_dtype': torch.float16, 'learnable_replace': True } global TRAIN_CONFIG def set_train_config(train_config): global TRAIN_CONFIG TRAIN_CONFIG = train_config def set_t5_encoder_path(path): global T5_ENCODER T5_ENCODER['MT5'] = path def create_logger(log_dir=None): logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) return logger def create_exp_folder(args, rank): if rank == 0: os.makedirs(args.results_dir, exist_ok=True) # print('args.results_dir', args.results_dir) _existed_experiments = os.listdir(args.results_dir) # print('_existed_experiments', _existed_experiments) existed_experiments = [] for x in _existed_experiments: if os.path.isdir(os.path.join(args.results_dir, x)): existed_experiments.append(x) if len(existed_experiments) == 0: experiment_index = 1 else: existed_experiments.sort() # print('existed_experiments', existed_experiments) experiment_index = max([int(x.split('-')[0]) for x in existed_experiments]) + 1 model_string_name = args.task_flag if args.task_flag else args.model.replace( "/", "-") # Create an experiment folder experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Stores saved model checkpoints checkpoint_dir = f"{experiment_dir}/checkpoints" if rank == 0: os.makedirs(checkpoint_dir, exist_ok=True) logger = create_logger(experiment_dir) logger.info(f"Experiment directory created at {experiment_dir}") else: logger = create_logger() experiment_dir = "" logger.info = print return experiment_dir, checkpoint_dir, logger def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir): def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"): cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}" if rank == 0: if args.use_fp16: model.module.save_pretrained(cur_ckpt_save_dir) else: model.save_pretrained(cur_ckpt_save_dir) checkpoint_path = "[Not rank 0. Disabled output.]" client_state = { "steps": train_steps, "epoch": epoch, "args": args } if ema is not None: client_state['ema'] = ema.state_dict() dst_paths = [] if train_steps % args.ckpt_every == 0: checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" try: if args.training_parts == "lora": save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt") else: model.save_checkpoint( checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt") dst_paths.append(checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") except: logger.error(f"Saved failed to {checkpoint_path}") if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps: save_name = "latest.pt" checkpoint_path = f"{checkpoint_dir}/{save_name}" try: if args.training_parts == "lora": save_lora_weight(checkpoint_dir, client_state, tag=f"{save_name}") else: model.save_checkpoint( checkpoint_dir, client_state=client_state, tag=f"{save_name}") dst_paths.append(checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") except: logger.error(f"Saved failed to {checkpoint_path}") return checkpoint_path @torch.no_grad() def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img): image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch # clip & mT5 text embedding text_embedding = text_embedding.to(device) text_embedding_mask = text_embedding_mask.to(device) encoder_hidden_states = text_encoder( text_embedding.to(device), attention_mask=text_embedding_mask.to(device), )[0] text_embedding_t5 = text_embedding_t5.to(device).squeeze(1) text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1) with torch.no_grad(): output_t5 = text_encoder_t5( input_ids=text_embedding_t5, attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None, output_hidden_states=True ) encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach( ) # additional condition image_meta_size = kwargs['image_meta_size'].to(device) style = kwargs['style'].to(device) if args.extra_fp16: image = image.half() image_meta_size = image_meta_size.half() if image_meta_size is not None else None # Map input images to latent space + normalize latents: image = image.to(device) vae_scaling_factor = vae.config.scaling_factor latents = vae.encode(image).latent_dist.sample().mul_( vae_scaling_factor).to(device) # positional embedding _, _, height, width = image.shape reso = f"{height}x{width}" # print("freqs_cis_img:", freqs_cis_img.keys()) # print("freqs_cis_img[reso]:", freqs_cis_img[reso]) cos_cis_img, sin_cis_img = freqs_cis_img[reso] # Model conditions model_kwargs = dict( encoder_hidden_states=encoder_hidden_states, text_embedding_mask=text_embedding_mask, encoder_hidden_states_t5=encoder_hidden_states_t5, text_embedding_mask_t5=text_embedding_mask_t5, image_meta_size=image_meta_size, style=style, cos_cis_img=cos_cis_img, sin_cis_img=sin_cis_img, ) return latents, model_kwargs def Core(args, LOG): if args.training_parts == "lora": args.use_ema = False assert torch.cuda.is_available(), "Training currently requires at least one GPU." batch_size = args.batch_size grad_accu_steps = args.grad_accu_steps global_batch_size = batch_size * grad_accu_steps seed = args.global_seed random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) print(f"Starting ,seed={seed}.") rank = 0 # Setup an experiment folder experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank) # Log all the arguments logger.info(sys.argv) logger.info(str(args)) # Save to a json file args_dict = vars(args) with open(f"{experiment_dir}/args.json", 'w') as f: json.dump(args_dict, f, indent=4) # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel." # If needed, just comment the following line. tf_logging.set_verbosity_error() # =========================================================================== # Building HYDIT # =========================================================================== logger.info("Building HYDIT Model.") # --------------------------------------------------------------------------- # Training sample base size, such as 256/512/1024. Notice that this size is # just a base size, not necessary the actual size of training samples. Actual # size of the training samples are correlated with `resolutions` when enabling # multi-resolution training. # --------------------------------------------------------------------------- image_size = args.image_size if type(image_size) == int: image_size = [image_size, image_size] if len(image_size) == 1: image_size = [image_size[0], image_size[0]] if len(image_size) != 2: raise ValueError(f"Invalid image size: {args.image_size}") assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \ f"got {image_size}" latent_size = [image_size[0] // 8, image_size[1] // 8] # initialize model by deepspeed model = HUNYUAN_DIT_MODELS[args.model](args, input_size=latent_size, log_fn=logger.info, ) # Multi-resolution / Single-resolution training. if args.multireso: resolutions = ResolutionGroup(image_size[0], align=16, step=args.reso_step, target_ratios=args.target_ratios).data else: resolutions = ResolutionGroup(image_size[0], align=16, target_ratios=['1:1']).data freqs_cis_img = init_image_posemb(args.rope_img, resolutions=resolutions, patch_size=model.patch_size, hidden_size=model.hidden_size, num_heads=model.num_heads, log_fn=logger.info, rope_real=args.rope_real, ) # Create EMA model and convert to fp16 if needed. ema = None device = "cuda" if args.use_ema: ema = EMA(args, model, device, logger) # Setup FP16 main model: if args.use_fp16: model = Float16Module(model, args) logger.info( f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}") diffusion = create_diffusion( noise_schedule=args.noise_schedule, predict_type=args.predict_type, learn_sigma=args.learn_sigma, mse_loss_weight_type=args.mse_loss_weight_type, beta_start=args.beta_start, beta_end=args.beta_end, noise_offset=args.noise_offset, ) # Setup VAE logger.info(f" Loading vae from {VAE_EMA_PATH}") if os.path.isdir(VAE_EMA_PATH): vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH) else: vae = AutoencoderKL.from_single_file(VAE_EMA_PATH) # Setup BERT text encoder logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}") text_encoder = BertModel.from_pretrained( TEXT_ENCODER, False, revision=None) # Setup BERT tokenizer: logger.info(f" Loading Bert tokenizer from {TOKENIZER}") tokenizer = BertTokenizer.from_pretrained(TOKENIZER) # Setup T5 text encoder from hydit.modules.text_encoder import MT5Embedder mt5_path = T5_ENCODER['MT5'] embedder_t5 = MT5Embedder( mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) tokenizer_t5 = embedder_t5.tokenizer text_encoder_t5 = embedder_t5.model if args.extra_fp16: logger.info(f" Using fp16 for extra modules: vae, text_encoder") vae = vae.half().to(device) text_encoder = text_encoder.half().to(device) text_encoder_t5 = text_encoder_t5.half().to(device) else: vae = vae.to(device) text_encoder = text_encoder.to(device) text_encoder_t5 = text_encoder_t5.to(device) logger.info( f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}") # =========================================================================== # Building Dataset # =========================================================================== logger.info(f"Building Streaming Dataset.") logger.info(f" Loading index file {args.index_file} (v2)") world_size = 1 dataset = TextImageArrowStream(args=args, resolution=image_size[0], random_flip=args.random_flip, log_fn=logger.info, index_file=args.index_file, multireso=args.multireso, batch_size=batch_size, world_size=world_size, random_shrink_size_cond=args.random_shrink_size_cond, merge_src_cond=args.merge_src_cond, uncond_p=args.uncond_p, text_ctx_len=args.text_len, tokenizer=tokenizer, uncond_p_t5=args.uncond_p_t5, text_ctx_len_t5=args.text_len_t5, tokenizer_t5=tokenizer_t5, ) if args.multireso: sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, shuffle=False, drop_last=True, batch_size=batch_size) else: sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, shuffle=False, drop_last=True) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler, num_workers=args.num_workers, pin_memory=True, drop_last=True) logger.info(f" Dataset contains {len(dataset):,} images.") logger.info(f" Index file: {args.index_file}.") if args.multireso: logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} ' f'and base size {dataset.index_manager.base_size}') logger.info(f'\n {dataset.index_manager.resolutions}') # =========================================================================== # Loading parameter # =========================================================================== logger.info(f"Loading parameter") start_epoch = 0 start_epoch_step = 0 train_steps = 0 # Resume checkpoint if needed model, ema, start_epoch, start_epoch_step, train_steps = model_resume( args, model, ema, logger) if args.training_parts == "lora": loraconfig = LoraConfig( r=args.rank, lora_alpha=args.rank, target_modules=args.target_modules ) if args.use_fp16: model.module = get_peft_model(model.module, loraconfig) else: model = get_peft_model(model, loraconfig) logger.info(f" Training parts: {args.training_parts}") # =========================================================================== # Training # =========================================================================== model.to(device) model.train() if args.use_ema: ema.eval() print(f" Worker {rank} ready.") iters_per_epoch = len(loader) logger.info( " ****************************** Running training ******************************") logger.info(f" Number GPUs: {world_size}") logger.info(f" Number training samples: {len(dataset):,}") logger.info( f" Number parameters: {sum(p.numel() for p in model.parameters()):,}") logger.info( f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}") logger.info( " ------------------------------------------------------------------------------") logger.info(f" Iters per epoch: {iters_per_epoch:,}") logger.info(f" Batch size per device: {batch_size}") logger.info( f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)") logger.info(f" Gradient Accu steps: {args.grad_accu_steps}") logger.info( f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}") logger.info( f" Training epochs: {start_epoch}/{args.epochs}") logger.info( f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}") logger.info( f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}") logger.info( " ------------------------------------------------------------------------------") logger.info(f" Noise schedule: {args.noise_schedule}") logger.info( f" Beta limits: ({args.beta_start}, {args.beta_end})") logger.info(f" Learn sigma: {args.learn_sigma}") logger.info(f" Prediction type: {args.predict_type}") logger.info(f" Noise offset: {args.noise_offset}") logger.info( " ------------------------------------------------------------------------------") logger.info( f" Using EMA model: {args.use_ema} ({args.ema_dtype})") if args.use_ema: logger.info( f" Using EMA decay: {ema.max_value if args.use_ema else None}") logger.info( f" Using EMA warmup power: {ema.power if args.use_ema else None}") logger.info(f" Using main model fp16: {args.use_fp16}") logger.info(f" Using extra modules fp16: {args.extra_fp16}") logger.info( " ------------------------------------------------------------------------------") logger.info(f" Experiment directory: {experiment_dir}") logger.info( " *******************************************************************************") if args.gc_interval > 0: gc.disable() gc.collect() # Variables for monitoring/logging purposes: log_steps = 0 running_loss = 0 start_time = time.time() if args.async_ema: ema_stream = torch.cuda.Stream() total_steps = args.epochs * len(loader) if args.use_fp16: optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=4e-6) else: optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay) easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, target_height=768, target_width=1280, train_steps=0) # torch.autograd.set_detect_anomaly(True) # with torch.autograd.detect_anomaly(): with tqdm(total=total_steps, initial=train_steps) as pbar: # Training loop for epoch in range(start_epoch, args.epochs): logger.info(f" Start random shuffle with seed={seed}") # Makesure all processors use the same seed to shuffle dataset. dataset.shuffle(seed=args.global_seed + epoch, fast=True) logger.info(f" End of random shuffle") # Move sampler to start_index if not args.multireso: start_index = start_epoch_step * world_size * batch_size if start_index != sampler.start_index: sampler.start_index = start_index # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning. start_epoch_step = 0 logger.info( f" Iters left this epoch: {len(loader):,}") logger.info(f" Beginning epoch {epoch}...") step = 0 mean_loss = 0 for batch in loader: step += 1 # Zero gradients, backward pass, and optimizer step: optimizer.zero_grad() latents, model_kwargs = prepare_model_inputs( args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img) # with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): if True: # training model by deepspeed while use fp16 if args.use_fp16: if args.use_ema and args.async_ema: with torch.cuda.stream(ema_stream): ema.update(model.module.module, step=step) torch.cuda.current_stream().wait_stream(ema_stream) loss_dict = diffusion.training_losses( model=model, x_start=latents, model_kwargs=model_kwargs) loss = loss_dict["loss"].mean() # print(f"step={step}, loss={loss.item()}") # model.backward(loss) loss.backward() # last_batch_iteration = ( # train_steps + 1) // (global_batch_size // (batch_size * world_size)) # model.step( # lr_kwargs={'last_batch_iteration': last_batch_iteration}) optimizer.step() if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1): if args.use_fp16: ema.update(model.module.module, step=step) else: ema.update(model.module, step=step) train_steps += 1 mean_loss += loss.item() pbar.update(1) pbar.set_description( f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") # =========================================================================== # Log loss values: # =========================================================================== # running_loss += loss.item() # log_steps += 1 # if train_steps % args.log_every == 0: # # Measure training speed: # torch.cuda.synchronize() # end_time = time.time() # steps_per_sec = log_steps / (end_time - start_time) # # Reduce loss history over all processes: # avg_loss = torch.tensor( # running_loss / log_steps, device=device) # # dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM) # # avg_loss = avg_loss.item() / world_size # # get lr from deepspeed fused optimizer # logger.info(f"(step={train_steps:07d}) " + # (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") + # f"Train Loss: {avg_loss:.4f}, " # # f"Lr: {opt.param_groups[0]['lr']:.6g}, " # f"Steps/Sec: {steps_per_sec:.2f}, " # f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}") # # Reset monitoring variables: # running_loss = 0 # log_steps = 0 # start_time = time.time() # collect gc: if args.gc_interval > 0 and (step % args.gc_interval == 0): gc.collect() if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps ) and train_steps > 0: logger.info( f" Saving checkpoint at step {train_steps}.") easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, target_height=768, target_width=1280, train_steps=train_steps) save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir) if train_steps >= args.max_training_steps: logger.info(f"Breaking step loop at {train_steps}.") break LOG({ "type": "sample_images", "global_step": train_steps, "total_steps": total_steps, # "latent": noise_pred_latent_path, }) if train_steps >= args.max_training_steps: logger.info(f"Breaking epoch loop at {epoch}.") break def model_resume(args, model, ema, logger): """ Load pretrained weights. """ start_epoch = 0 start_epoch_step = 0 train_steps = 0 resume_path = UNET_PATH if not Path(resume_path).exists(): raise FileNotFoundError( f" Cannot find checkpoint from {resume_path}") logger.info(f"Resume from checkpoint {resume_path}") if args.resume_split: # Resume main model resume_ckpt_module = torch.load( resume_path, map_location=lambda storage, loc: storage) model.load_state_dict(resume_ckpt_module, strict=False) # Resume ema model if args.use_ema: if args.module_to_ema: if "resume_ckpt_module" in locals(): logger.info(f" Resume ema model from main states.") ema.load_state_dict(resume_ckpt_module, strict=args.strict) else: logger.info(f" Resume ema model from module states.") resume_ckpt_module = torch.load( resume_path, map_location=lambda storage, loc: storage) ema.load_state_dict(resume_ckpt_module, strict=args.strict) else: if "resume_ckpt_ema" in locals(): logger.info(f" Resume ema model from EMA states.") ema.load_state_dict(resume_ckpt_ema, strict=args.strict) else: logger.info(f" Resume ema model from EMA states.") resume_ckpt_ema = torch.load(resume_path, map_location=lambda storage, loc: storage) ema.load_state_dict(resume_ckpt_ema, strict=args.strict) else: raise ValueError( " “If `resume` is True, then either `resume_split` must be true.”") return model, ema, start_epoch, start_epoch_step, train_steps def easy_sample_images( args, vae=None, text_encoder=None, tokenizer=None, model=None, embedder_t5=None, target_height=768, target_width=1280, prompt="A photo of a girl with a hat on a sunny day", negative_prompt="", batch_size=1, guidance_scale=5.0, infer_steps=20, sampler='ddpm', train_steps=0, ): with torch.no_grad(): from hydit.diffusion.pipeline import StableDiffusionPipeline from diffusers import schedulers from hydit.constants import SAMPLER_FACTORY from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed from hydit.modules.models import HUNYUAN_DIT_CONFIG workspace_dir = TRAIN_CONFIG.get("workspace_dir") sample_config_file = TRAIN_CONFIG.get("sample_config_file", None) if sample_config_file is None: print("sample_config_file is not set.") return try: sample_config = json.load(open(sample_config_file, "r")) except Exception as e: print(f"Failed to load sample_config_file: {sample_config_file}") return sample_images_dir = os.path.join(workspace_dir, "sample_images") os.makedirs(sample_images_dir, exist_ok=True) # Load sampler from factory kwargs = SAMPLER_FACTORY[sampler]['kwargs'] scheduler = SAMPLER_FACTORY[sampler]['scheduler'] # Build scheduler according to the sampler. scheduler_class = getattr(schedulers, scheduler) scheduler = scheduler_class(**kwargs) # Set timesteps for inference steps. scheduler.set_timesteps(infer_steps, "cuda") def calc_rope(height, width): model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] patch_size = model_config['patch_size'] head_size = model_config['hidden_size'] // model_config['num_heads'] th = height // 8 // patch_size tw = width // 8 // patch_size base_size = 512 // 8 // patch_size start, stop = get_fill_resize_and_crop((th, tw), base_size) sub_args = [start, stop, (th, tw)] rope = get_2d_rotary_pos_embed(head_size, *sub_args) return rope pipeline = StableDiffusionPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=model, scheduler=scheduler, feature_extractor=None, safety_checker=None, requires_safety_checker=False, embedder_t5=embedder_t5, ) pipeline.to("cuda") style = torch.as_tensor([0, 0] * batch_size, device="cuda") src_size_cond = (target_width, target_height) size_cond = list(src_size_cond) + [target_width, target_height, 0, 0] image_meta_size = torch.as_tensor( [size_cond] * 2 * batch_size, device="cuda",) if type(sample_config) != list: sample_config = [sample_config] for i, sample in enumerate(sample_config): prompt = sample.get("prompt", "") negative_prompt = sample.get("negative_prompt", "") guidance_scale = sample.get("cfg", guidance_scale) infer_steps = sample.get("steps", infer_steps) width = sample.get("width", target_width) height = sample.get("height", target_height) freqs_cis_img = calc_rope(height, width) samples = pipeline( height=height, width=width, prompt=prompt, negative_prompt=negative_prompt, num_images_per_prompt=batch_size, guidance_scale=guidance_scale, num_inference_steps=infer_steps, style=style, return_dict=False, use_fp16=True, learn_sigma=args.learn_sigma, freqs_cis_img=freqs_cis_img, image_meta_size=image_meta_size, )[0] # print("samples:",type(samples),) # input("Press Enter to continue...") # print("samples:",samples,) if type(samples) == list: pil_image = samples[0] else: pil_image = samples sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png" sample_filename_path = os.path.join( sample_images_dir, sample_filename) pil_image.save(sample_filename_path) return None