Guest User

TI_optimization

a guest
Oct 30th, 2022
1,736
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Diff 3.97 KB | Source Code | 0 0
  1. diff --git a/modules/shared.py b/modules/shared.py
  2. index fb84afd..900f920 100644
  3. --- a/modules/shared.py
  4. +++ b/modules/shared.py
  5. @@ -261,6 +261,7 @@ options_templates.update(options_section(('training', "Training"), {
  6.      "dataset_filename_join_string": OptionInfo(" ", "Filename join string"),
  7.      "training_image_repeats_per_epoch": OptionInfo(1, "Number of repeats for a single input image per epoch; used only for displaying epoch number", gr.Number, {"precision": 0}),
  8.      "training_write_csv_every": OptionInfo(500, "Save an csv containing the loss to log directory every N steps, 0 to disable"),
  9. +    "training_xattention_optimizations": OptionInfo(False, "Use cross attention optimizations while training"),
  10.  }))
  11.  
  12.  options_templates.update(options_section(('sd', "Stable Diffusion"), {
  13. diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
  14. index 17dfb22..b0a1d26 100644
  15. --- a/modules/textual_inversion/textual_inversion.py
  16. +++ b/modules/textual_inversion/textual_inversion.py
  17. @@ -214,6 +214,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
  18.      filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')
  19.  
  20.      log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), embedding_name)
  21. +    unload = shared.opts.unload_models_when_training
  22.  
  23.      if save_embedding_every > 0:
  24.          embedding_dir = os.path.join(log_directory, "embeddings")
  25. @@ -238,6 +239,8 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
  26.      shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
  27.      with torch.autocast("cuda"):
  28.          ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, device=devices.device, template_file=template_file, batch_size=batch_size)
  29. +    if unload:
  30. +        shared.sd_model.first_stage_model.to(devices.cpu)
  31.  
  32.      hijack = sd_hijack.model_hijack
  33.  
  34. @@ -303,6 +306,9 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
  35.          if images_dir is not None and steps_done % create_image_every == 0:
  36.              forced_filename = f'{embedding_name}-{steps_done}'
  37.              last_saved_image = os.path.join(images_dir, forced_filename)
  38. +
  39. +            shared.sd_model.first_stage_model.to(devices.device)
  40. +
  41.              p = processing.StableDiffusionProcessingTxt2Img(
  42.                  sd_model=shared.sd_model,
  43.                  do_not_save_grid=True,
  44. @@ -330,6 +336,9 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
  45.              processed = processing.process_images(p)
  46.              image = processed.images[0]
  47.  
  48. +            if unload:
  49. +                shared.sd_model.first_stage_model.to(devices.cpu)
  50. +
  51.              shared.state.current_image = image
  52.  
  53.              if save_image_with_stored_embedding and os.path.exists(last_saved_file) and embedding_yet_to_be_embedded:
  54. diff --git a/modules/textual_inversion/ui.py b/modules/textual_inversion/ui.py
  55. index e712284..d679e6f 100644
  56. --- a/modules/textual_inversion/ui.py
  57. +++ b/modules/textual_inversion/ui.py
  58. @@ -25,8 +25,10 @@ def train_embedding(*args):
  59.  
  60.      assert not shared.cmd_opts.lowvram, 'Training models with lowvram not possible'
  61.  
  62. +    apply_optimizations = shared.opts.training_xattention_optimizations
  63.      try:
  64. -        sd_hijack.undo_optimizations()
  65. +        if not apply_optimizations:
  66. +            sd_hijack.undo_optimizations()
  67.  
  68.          embedding, filename = modules.textual_inversion.textual_inversion.train_embedding(*args)
  69.  
  70. @@ -38,5 +40,6 @@ Embedding saved to {html.escape(filename)}
  71.      except Exception:
  72.          raise
  73.      finally:
  74. -        sd_hijack.apply_optimizations()
  75. +        if not apply_optimizations:
  76. +            sd_hijack.apply_optimizations()
  77.  
Add Comment
Please, Sign In to add comment