Update utils.py

utils.py

@@ -10,8 +10,6 @@ from scipy.io.wavfile import read
 import torch
 import regex as re
 
-import loralib as lora
-
 MATPLOTLIB_FLAG = False
 
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
@@ -146,91 +144,6 @@ def tag_cke(text,prev_sentence=None):
     return prev_lang,tagged_text
 
 
-def load_lora_checkpoint(checkpoint_path, model, optimizer=None, generator_path = "./pretrained_models/G_latest.pth"):
-  assert os.path.isfile(checkpoint_path)
-  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-  iteration = checkpoint_dict['iteration']
-  learning_rate = checkpoint_dict['learning_rate']
-  if optimizer is not None:
-    optimizer.load_state_dict(checkpoint_dict['optimizer'])
-#   saved_state_dict = checkpoint_dict['model']
-  generator_state_dict = torch.load(generator_path)['model']
-  lora_state_dict = checkpoint_dict['model']
-  new_state_dict = {}
-  for k, v in lora_state_dict.items():
-    try:
-        if k == 'emb_g.weight':
-            if drop_speaker_emb:
-                new_state_dict[k] = v
-                continue
-            v[:lora_state_dict[k].shape[0], :] = lora_state_dict[k]
-            new_state_dict[k] = v
-        else:
-            new_state_dict[k] = lora_state_dict[k]
-    except:
-        logger.info("%s is not in the checkpoint" % k)
-        new_state_dict[k] = v
-  if hasattr(model, 'module'):
-    model.module.load_state_dict(generator_state_dict, strict=False)
-    model.module.load_state_dict(new_state_dict, strict=False)
-    # lora.mark_only_lora_as_trainable(model.module)
-  else:
-    model.load_state_dict(generator_state_dict, strict=False)
-    model.load_state_dict(new_state_dict, strict=False)
-    # lora.mark_only_lora_as_trainable(model)
-  logger.info("Loaded checkpoint '{}' (iteration {})" .format(
-    checkpoint_path, iteration))
-
-  return model, optimizer, learning_rate, iteration
-
-def save_lora_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
-  logger.info("Saving model and optimizer state at iteration {} to {}".format(
-    iteration, checkpoint_path))
-  if hasattr(model, 'module'):
-    state_dict = lora.lora_state_dict(model.module)
-  else:
-    state_dict = lora.lora_state_dict(model)
-  torch.save({'model': state_dict,
-              'iteration': iteration,
-              'optimizer': optimizer.state_dict() if optimizer is not None else None,
-              'learning_rate': learning_rate}, checkpoint_path)
-
-def load_lora_checkpoint_fix(checkpoint_path, model, optimizer=None, drop_speaker_emb=False):
-    assert os.path.isfile(checkpoint_path)
-    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-    iteration = checkpoint_dict['iteration']
-    learning_rate = checkpoint_dict['learning_rate']
-    if optimizer is not None:
-        optimizer.load_state_dict(checkpoint_dict['optimizer'])
-    saved_state_dict = checkpoint_dict['model']
-    if hasattr(model, 'module'):
-        state_dict = model.module.state_dict()
-    else:
-        state_dict = model.state_dict()
-    new_state_dict = {}
-    for k, v in state_dict.items():
-        try:
-            if k == 'emb_g.weight':
-                if drop_speaker_emb:
-                    new_state_dict[k] = v
-                    continue
-                v[:saved_state_dict[k].shape[0], :] = saved_state_dict[k]
-                new_state_dict[k] = v
-            else:
-                new_state_dict[k] = saved_state_dict[k]
-        except:
-            logger.info("%s is not in the checkpoint" % k)
-            new_state_dict[k] = v
-    if hasattr(model, 'module'):
-        model.module.load_state_dict(new_state_dict, strict=False)
-        lora.mark_only_lora_as_trainable(model.module)
-    else:
-        model.load_state_dict(new_state_dict, strict=False)
-        lora.mark_only_lora_as_trainable(model)
-    logger.info("Loaded checkpoint '{}' (iteration {})".format(
-        checkpoint_path, iteration))
-    return model, optimizer, learning_rate, iteration
-
 def load_checkpoint(checkpoint_path, model, optimizer=None, drop_speaker_emb=False):
     assert os.path.isfile(checkpoint_path)
     checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')