From 5f4e14eac9c5b766e4931078c69aa065fe8a7616 Mon Sep 17 00:00:00 2001 From: wailovet Date: Fri, 28 Jun 2024 01:44:31 +0800 Subject: [PATCH] 'update' --- hook_HYDiT_main.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py index d6e3292..b65064b 100644 --- a/hook_HYDiT_main.py +++ b/hook_HYDiT_main.py @@ -576,12 +576,12 @@ def Core(args, LOG): with torch.cuda.stream(ema_stream): ema.update(model.module.module, step=step) torch.cuda.current_stream().wait_stream(ema_stream) - - with torch.cuda.amp.autocast(enabled=args.use_fp16): + + with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): loss_dict = diffusion.training_losses( model=model, x_start=latents, model_kwargs=model_kwargs) - loss = loss_dict["loss"].mean() + loss = loss_dict["loss"].mean() # print(f"step={step}, loss={loss.item()}") @@ -641,7 +641,7 @@ def Core(args, LOG): 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: + ) and train_steps > 0: logger.info( f" Saving checkpoint at step {train_steps}.") easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, @@ -718,6 +718,8 @@ 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 + + def easy_sample_images( args, vae=None, @@ -737,7 +739,6 @@ def easy_sample_images( ): with torch.no_grad(): - workspace_dir = TRAIN_CONFIG.get("workspace_dir") sample_config_file = TRAIN_CONFIG.get("sample_config_file", None) if sample_config_file is None: @@ -804,7 +805,7 @@ def easy_sample_images( height = sample.get("height", target_height) freqs_cis_img = calc_rope(height, width) - + try: samples = pipeline( height=height,