Huge memory optimizations, many big fixes

jobs/process/BaseExtensionProcess.py

@@ -4,10 +4,6 @@ from jobs.process.BaseProcess import BaseProcess
 
 
 class BaseExtensionProcess(BaseProcess):
-    process_id: int
-    config: OrderedDict
-    progress_bar: ForwardRef('tqdm') = None
-
     def __init__(
             self,
             process_id: int,
@@ -15,6 +11,9 @@ class BaseExtensionProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
+        self.process_id: int
+        self.config: OrderedDict
+        self.progress_bar: ForwardRef('tqdm') = None
 
     def run(self):
         super().run()

jobs/process/BaseExtractProcess.py

@@ -12,11 +12,6 @@ from toolkit.train_tools import get_torch_dtype
 
 
 class BaseExtractProcess(BaseProcess):
-    process_id: int
-    config: OrderedDict
-    output_folder: str
-    output_filename: str
-    output_path: str
 
     def __init__(
             self,
@@ -25,6 +20,10 @@ class BaseExtractProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
+        self.config: OrderedDict
+        self.output_folder: str
+        self.output_filename: str
+        self.output_path: str
         self.process_id = process_id
         self.job = job
         self.config = config

jobs/process/BaseMergeProcess.py

@@ -9,8 +9,6 @@ from toolkit.train_tools import get_torch_dtype
 
 
 class BaseMergeProcess(BaseProcess):
-    process_id: int
-    config: OrderedDict
 
     def __init__(
             self,
@@ -19,6 +17,8 @@ class BaseMergeProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
+        self.process_id: int
+        self.config: OrderedDict
         self.output_path = self.get_conf('output_path', required=True)
         self.dtype = self.get_conf('dtype', self.job.dtype)
         self.torch_dtype = get_torch_dtype(self.dtype)

jobs/process/BaseProcess.py

@@ -4,7 +4,6 @@ from collections import OrderedDict
 
 
 class BaseProcess(object):
-    meta: OrderedDict
 
     def __init__(
             self,
@@ -13,6 +12,7 @@ class BaseProcess(object):
             config: OrderedDict
     ):
         self.process_id = process_id
+        self.meta: OrderedDict
         self.job = job
         self.config = config
         self.raw_process_config = config

jobs/process/BaseSDTrainProcess.py

@@ -9,6 +9,7 @@ from toolkit.data_loader import get_dataloader_from_datasets
 from toolkit.embedding import Embedding
 from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.optimizer import get_optimizer
+from toolkit.paths import CONFIG_ROOT
 
 from toolkit.scheduler import get_lr_scheduler
 from toolkit.stable_diffusion_model import StableDiffusion
@@ -31,11 +32,12 @@ def flush():
 
 
 class BaseSDTrainProcess(BaseTrainProcess):
-    sd: StableDiffusion
-    embedding: Union[Embedding, None] = None
 
     def __init__(self, process_id: int, job, config: OrderedDict, custom_pipeline=None):
         super().__init__(process_id, job, config)
+        self.sd: StableDiffusion
+        self.embedding: Union[Embedding, None] = None
+
         self.custom_pipeline = custom_pipeline
         self.step_num = 0
         self.start_step = 0
@@ -344,7 +346,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
 
             # remove grads for these
             noisy_latents.requires_grad = False
+            noisy_latents = noisy_latents.detach()
             noise.requires_grad = False
+            noise = noise.detach()
 
         return noisy_latents, noise, timesteps, conditioned_prompts, imgs
 

jobs/process/BaseTrainProcess.py

@@ -14,11 +14,6 @@ if TYPE_CHECKING:
 
 
 class BaseTrainProcess(BaseProcess):
-    process_id: int
-    config: OrderedDict
-    writer: 'SummaryWriter'
-    job: Union['TrainJob', 'BaseJob', 'ExtensionJob']
-    progress_bar: 'tqdm' = None
 
     def __init__(
             self,
@@ -27,6 +22,12 @@ class BaseTrainProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
+        self.process_id: int
+        self.config: OrderedDict
+        self.writer: 'SummaryWriter'
+        self.job: Union['TrainJob', 'BaseJob', 'ExtensionJob']
+        self.progress_bar: 'tqdm' = None
+
         self.progress_bar = None
         self.writer = None
         self.training_folder = self.get_conf('training_folder', self.job.training_folder if hasattr(self.job, 'training_folder') else None)

jobs/process/GenerateProcess.py

@@ -16,9 +16,9 @@ import random
 
 
 class GenerateConfig:
-    prompts: List[str]
 
     def __init__(self, **kwargs):
+        self.prompts: List[str]
         self.sampler = kwargs.get('sampler', 'ddpm')
         self.width = kwargs.get('width', 512)
         self.height = kwargs.get('height', 512)

jobs/process/ModRescaleLoraProcess.py

@@ -24,6 +24,9 @@ class ModRescaleLoraProcess(BaseProcess):
             config: OrderedDict
     ):
         super().__init__(process_id, job, config)
+        self.process_id: int
+        self.config: OrderedDict
+        self.progress_bar: ForwardRef('tqdm') = None
         self.input_path = self.get_conf('input_path', required=True)
         self.output_path = self.get_conf('output_path', required=True)
         self.replace_meta = self.get_conf('replace_meta', default=False)