Add cleanup after training

1 files changed, 105 insertions, 96 deletions

diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 38f35c58..170d5ea4 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -398,110 +398,112 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
     forced_filename = "<none>"

 

     pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)

-    for i, entries in pbar:

-        hypernetwork.step = i + ititial_step

-        if len(loss_dict) > 0:

-            previous_mean_losses = [i[-1] for i in loss_dict.values()]

-            previous_mean_loss = mean(previous_mean_losses)

-            

-        scheduler.apply(optimizer, hypernetwork.step)

-        if scheduler.finished:

-            break

-

-        if shared.state.interrupted:

-            break

-

-        with torch.autocast("cuda"):

-            c = stack_conds([entry.cond for entry in entries]).to(devices.device)

-            # c = torch.vstack([entry.cond for entry in entries]).to(devices.device)

-            x = torch.stack([entry.latent for entry in entries]).to(devices.device)

-            loss = shared.sd_model(x, c)[0]

-            del x

-            del c

-

-            losses[hypernetwork.step % losses.shape[0]] = loss.item()

-            for entry in entries:

-                loss_dict[entry.filename].append(loss.item())

-                

-            optimizer.zero_grad()

-            weights[0].grad = None

-            loss.backward()

 

-            if weights[0].grad is None:

-                steps_without_grad += 1

+    try:

+        for i, entries in pbar:

+            hypernetwork.step = i + ititial_step

+            if len(loss_dict) > 0:

+                previous_mean_losses = [i[-1] for i in loss_dict.values()]

+                previous_mean_loss = mean(previous_mean_losses)

+                

+            scheduler.apply(optimizer, hypernetwork.step)

+            if scheduler.finished:

+                break

+

+            if shared.state.interrupted:

+                break

+

+            with torch.autocast("cuda"):

+                c = stack_conds([entry.cond for entry in entries]).to(devices.device)

+                # c = torch.vstack([entry.cond for entry in entries]).to(devices.device)

+                x = torch.stack([entry.latent for entry in entries]).to(devices.device)

+                loss = shared.sd_model(x, c)[0]

+                del x

+                del c

+

+                losses[hypernetwork.step % losses.shape[0]] = loss.item()

+                for entry in entries:

+                    loss_dict[entry.filename].append(loss.item())

+                    

+                optimizer.zero_grad()

+                weights[0].grad = None

+                loss.backward()

+

+                if weights[0].grad is None:

+                    steps_without_grad += 1

+                else:

+                    steps_without_grad = 0

+                assert steps_without_grad < 10, 'no gradient found for the trained weight after backward() for 10 steps in a row; this is a bug; training cannot continue'

+

+                optimizer.step()

+

+            steps_done = hypernetwork.step + 1

+

+            if torch.isnan(losses[hypernetwork.step % losses.shape[0]]): 

+                raise RuntimeError("Loss diverged.")

+            

+            if len(previous_mean_losses) > 1:

+                std = stdev(previous_mean_losses)

             else:

-                steps_without_grad = 0

-            assert steps_without_grad < 10, 'no gradient found for the trained weight after backward() for 10 steps in a row; this is a bug; training cannot continue'

-

-            optimizer.step()

-

-        steps_done = hypernetwork.step + 1

-

-        if torch.isnan(losses[hypernetwork.step % losses.shape[0]]): 

-            raise RuntimeError("Loss diverged.")

-        

-        if len(previous_mean_losses) > 1:

-            std = stdev(previous_mean_losses)

-        else:

-            std = 0

-        dataset_loss_info = f"dataset loss:{mean(previous_mean_losses):.3f}" + u"\u00B1" + f"({std / (len(previous_mean_losses) ** 0.5):.3f})"

-        pbar.set_description(dataset_loss_info)

-

-        if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:

-            # Before saving, change name to match current checkpoint.

-            hypernetwork.name = f'{hypernetwork_name}-{steps_done}'

-            last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork.name}.pt')

-            hypernetwork.save(last_saved_file)

-

-        textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, len(ds), {

-            "loss": f"{previous_mean_loss:.7f}",

-            "learn_rate": scheduler.learn_rate

-        })

-

-        if images_dir is not None and steps_done % create_image_every == 0:

-            forced_filename = f'{hypernetwork_name}-{steps_done}'

-            last_saved_image = os.path.join(images_dir, forced_filename)

-

-            optimizer.zero_grad()

-            shared.sd_model.cond_stage_model.to(devices.device)

-            shared.sd_model.first_stage_model.to(devices.device)

-

-            p = processing.StableDiffusionProcessingTxt2Img(

-                sd_model=shared.sd_model,

-                do_not_save_grid=True,

-                do_not_save_samples=True,

-            )

+                std = 0

+            dataset_loss_info = f"dataset loss:{mean(previous_mean_losses):.3f}" + u"\u00B1" + f"({std / (len(previous_mean_losses) ** 0.5):.3f})"

+            pbar.set_description(dataset_loss_info)

+

+            if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:

+                # Before saving, change name to match current checkpoint.

+                hypernetwork.name = f'{hypernetwork_name}-{steps_done}'

+                last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork.name}.pt')

+                hypernetwork.save(last_saved_file)

+

+            textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, len(ds), {

+                "loss": f"{previous_mean_loss:.7f}",

+                "learn_rate": scheduler.learn_rate

+            })

+

+            if images_dir is not None and steps_done % create_image_every == 0:

+                forced_filename = f'{hypernetwork_name}-{steps_done}'

+                last_saved_image = os.path.join(images_dir, forced_filename)

+

+                optimizer.zero_grad()

+                shared.sd_model.cond_stage_model.to(devices.device)

+                shared.sd_model.first_stage_model.to(devices.device)

+

+                p = processing.StableDiffusionProcessingTxt2Img(

+                    sd_model=shared.sd_model,

+                    do_not_save_grid=True,

+                    do_not_save_samples=True,

+                )

 

-            if preview_from_txt2img:

-                p.prompt = preview_prompt

-                p.negative_prompt = preview_negative_prompt

-                p.steps = preview_steps

-                p.sampler_index = preview_sampler_index

-                p.cfg_scale = preview_cfg_scale

-                p.seed = preview_seed

-                p.width = preview_width

-                p.height = preview_height

-            else:

-                p.prompt = entries[0].cond_text

-                p.steps = 20

+                if preview_from_txt2img:

+                    p.prompt = preview_prompt

+                    p.negative_prompt = preview_negative_prompt

+                    p.steps = preview_steps

+                    p.sampler_index = preview_sampler_index

+                    p.cfg_scale = preview_cfg_scale

+                    p.seed = preview_seed

+                    p.width = preview_width

+                    p.height = preview_height

+                else:

+                    p.prompt = entries[0].cond_text

+                    p.steps = 20

 

-            preview_text = p.prompt

+                preview_text = p.prompt

 

-            processed = processing.process_images(p)

-            image = processed.images[0] if len(processed.images)>0 else None

+                processed = processing.process_images(p)

+                image = processed.images[0] if len(processed.images)>0 else None

 

-            if unload:

-                shared.sd_model.cond_stage_model.to(devices.cpu)

-                shared.sd_model.first_stage_model.to(devices.cpu)

+                if unload:

+                    shared.sd_model.cond_stage_model.to(devices.cpu)

+                    shared.sd_model.first_stage_model.to(devices.cpu)

 

-            if image is not None:

-                shared.state.current_image = image

-                last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)

-                last_saved_image += f", prompt: {preview_text}"

+                if image is not None:

+                    shared.state.current_image = image

+                    last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)

+                    last_saved_image += f", prompt: {preview_text}"

 

-        shared.state.job_no = hypernetwork.step

+            shared.state.job_no = hypernetwork.step

 

-        shared.state.textinfo = f"""

+            shared.state.textinfo = f"""

 <p>

 Loss: {previous_mean_loss:.7f}<br/>

 Step: {hypernetwork.step}<br/>

@@ -510,7 +512,14 @@ Last saved hypernetwork: {html.escape(last_saved_file)}<br/>
 Last saved image: {html.escape(last_saved_image)}<br/>

 </p>

 """

-        

+    finally:

+        if weights:

+            for weight in weights:

+                weight.requires_grad = False

+        if unload:

+            shared.sd_model.cond_stage_model.to(devices.device)

+            shared.sd_model.first_stage_model.to(devices.device)

+

     report_statistics(loss_dict)

     checkpoint = sd_models.select_checkpoint()