unified move-model

CHANGELOG.md

@@ -217,6 +217,7 @@ As of this release, default backend is set to **diffusers** as its more feature
   - **model load to gpu**  
     new option in settings->diffusers allowing models to be loaded directly to GPU while keeping RAM free  
     this option is not compatible with any kind of model offloading as model is expected to stay in GPU  
+    additionally, all model-moves can now be traced with env variable `SD_MOVE_DEBUG`  
   - **xyz grid**
     - range control  
       example: `5.0-6.0:3` will generate 3 images with values `5.0,5.5,6.0`  

modules/control/run.py

@@ -245,9 +245,8 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
 
     original_pipeline = shared.sd_model
     shared.sd_model = pipe
-    if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-        shared.sd_model.to(shared.device)
-    shared.sd_model.to(device=devices.device, dtype=devices.dtype)
+    sd_models.move_model(shared.sd_model, shared.device)
+    shared.sd_model.to(dtype=devices.dtype)
     debug(f'Control device={devices.device} dtype={devices.dtype}')
     sd_models.copy_diffuser_options(shared.sd_model, original_pipeline) # copy options from original pipeline
     sd_models.set_diffuser_options(shared.sd_model)

modules/face/instantid.py

@@ -57,8 +57,8 @@ def instant_id(p: processing.StableDiffusionProcessing, app, source_image, stren
     sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc.
     shared.sd_model.load_ip_adapter_instantid(face_adapter, scale=strength)
     shared.sd_model.set_ip_adapter_scale(strength)
-    if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-        shared.sd_model.to(shared.device, devices.dtype) # move pipeline if needed, but don't touch if its under automatic managment
+    sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
+    shared.sd_model.to(dtype=devices.dtype)
 
     # pipeline specific args
     orig_prompt_attention = shared.opts.prompt_attention

modules/face/photomaker.py

@@ -1,6 +1,6 @@
 import os
 import huggingface_hub as hf
-from modules import shared, processing, sd_models
+from modules import shared, processing, sd_models, devices
 
 
 def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger, strength, start): # pylint: disable=arguments-differ
@@ -42,8 +42,8 @@ def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger,
     )
     sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) # copy options from original pipeline
     sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc.
-    if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-        shared.sd_model.to(shared.device) # move pipeline if needed, but don't touch if its under automatic managment
+    sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
+    shared.sd_model.to(dtype=devices.dtype)
 
     orig_prompt_attention = shared.opts.prompt_attention
     shared.opts.data['prompt_attention'] = 'Fixed attention' # otherwise need to deal with class_tokens_mask

modules/interrogate.py

@@ -164,7 +164,7 @@ class InterrogateModels:
         res = ""
         shared.state.begin('interrogate')
         try:
-            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+            if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
                 lowvram.send_everything_to_cpu()
                 devices.torch_gc()
             self.load()

modules/processing.py

@@ -302,7 +302,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if not shared.opts.keep_incomplete and shared.state.interrupted:
                 x_samples_ddim = []
 
-            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram and shared.backend == shared.Backend.ORIGINAL:
+            if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
                 lowvram.send_everything_to_cpu()
                 devices.torch_gc()
             if p.scripts is not None and isinstance(p.scripts, scripts.ScriptRunner):

modules/processing_diffusers.py

@@ -371,8 +371,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         shared.sd_model = orig_pipeline
         return results
 
-    if shared.opts.diffusers_move_base and not getattr(shared.sd_model, 'has_accelerate', False):
-        shared.sd_model.to(devices.device)
+    if shared.opts.diffusers_move_base:
+        sd_models.move_model(shared.sd_model, devices.device)
 
     # recompile if a paramater chages
     recompile_model()
@@ -505,15 +505,14 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         shared.state.job_count +=1
         if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
             save_intermediate(latents=output.images, suffix="-before-refiner")
-        if shared.opts.diffusers_move_base and not getattr(shared.sd_model, 'has_accelerate', False):
+        if shared.opts.diffusers_move_base:
             shared.log.debug('Moving to CPU: model=base')
-            shared.sd_model.to(devices.cpu)
-            devices.torch_gc()
+            sd_models.move_model(shared.sd_model, devices.cpu)
         if shared.state.interrupted or shared.state.skipped:
             shared.sd_model = orig_pipeline
             return results
-        if shared.opts.diffusers_move_refiner and not getattr(shared.sd_refiner, 'has_accelerate', False):
-            shared.sd_refiner.to(devices.device)
+        if shared.opts.diffusers_move_refiner:
+            sd_models.move_model(shared.sd_refiner, devices.device)
         p.ops.append('refine')
         p.is_refiner_pass = True
         shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
@@ -558,10 +557,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 for refiner_image in refiner_images:
                     results.append(refiner_image)
 
-        if shared.opts.diffusers_move_refiner and not getattr(shared.sd_refiner, 'has_accelerate', False):
+        if shared.opts.diffusers_move_refiner:
             shared.log.debug('Moving to CPU: model=refiner')
-            shared.sd_refiner.to(devices.cpu)
-            devices.torch_gc()
+            sd_models.move_model(shared.sd_refiner, devices.cpu)
         shared.state.job = prev_job
         shared.state.nextjob()
         p.is_refiner_pass = False

modules/processing_vae.py

@@ -36,10 +36,9 @@ def full_vae_decode(latents, model):
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
         shared.log.debug('Moving to CPU: model=UNet')
         unet_device = model.unet.device
-        model.unet.to(devices.cpu)
-        devices.torch_gc()
+        sd_models.move_model(model.unet, devices.cpu)
     if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'):
-        model.vae.to(devices.device)
+        sd_models.move_model(model.vae, devices.device)
     latents.to(model.vae.device)
 
     upcast = (model.vae.dtype == torch.float16) and getattr(model.vae.config, 'force_upcast', False) and hasattr(model, 'upcast_vae')
@@ -57,7 +56,7 @@ def full_vae_decode(latents, model):
             devices.torch_gc(force=True)
 
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
-        model.unet.to(unet_device)
+        sd_models.move_model(model.unet, unet_device)
     t1 = time.time()
     debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}')
     return decoded
@@ -68,13 +67,12 @@ def full_vae_encode(image, model):
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
         debug('Moving to CPU: model=UNet')
         unet_device = model.unet.device
-        model.unet.to(devices.cpu)
-        devices.torch_gc()
+        sd_models.move_model(model.unet, devices.cpu)
     if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'):
-        model.vae.to(devices.device)
+        sd_models.move_model(model.vae, devices.device)
     encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample()
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):
-        model.unet.to(unet_device)
+        sd_models.move_model(model.unet, unet_device)
     return encoded
 
 

modules/sd_models.py

@@ -37,6 +37,7 @@ sd_metadata_file = os.path.join(paths.data_path, "metadata.json")
 sd_metadata = None
 sd_metadata_pending = 0
 sd_metadata_timer = 0
+debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
 class CheckpointInfo:
@@ -723,6 +724,16 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
         sd_model.unet.to(memory_format=torch.channels_last)
 
 
+def move_model(model, device=None):
+    if model is not None and not getattr(model, 'has_accelerate', False):
+        try:
+            model.to(device)
+            debug_move(f'Model move: to={device} class={model.__class__} function={sys._getframe(1).f_code.co_name}') # pylint: disable=protected-access
+        except Exception as e:
+            shared.log.error(f'Model move: to={device} {e}')
+        devices.torch_gc()
+
+
 def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model'): # pylint: disable=unused-argument
     if shared.cmd_opts.profile:
         import cProfile
@@ -910,7 +921,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                 free_vram = gpu_vram.get('total', 0) - gpu_vram.get('used', 0)
                 refiner_enough_vram = free_vram >= 7 if "StableDiffusionXL" in sd_model.__class__.__name__ else 3
                 if not shared.opts.diffusers_move_base and refiner_enough_vram:
-                    sd_model.to(devices.device)
+                    move_model(sd_model, devices.device)
                     base_sent_to_cpu=False
                 else:
                     if not refiner_enough_vram and not (shared.opts.diffusers_move_base and shared.opts.diffusers_move_refiner):
@@ -921,14 +932,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                             shared.opts.diffusers_move_base=True
                             shared.opts.diffusers_move_refiner=True
                     shared.log.debug('Moving base model to CPU')
-                    if model_data.sd_model is not None:
-                        model_data.sd_model.to(devices.cpu)
+                    move_model(model_data.sd_model, devices.cpu)
                     devices.torch_gc(force=True)
-                    sd_model.to(devices.device)
+                    move_model(sd_model, devices.device)
                     base_sent_to_cpu=True
-            elif not getattr(sd_model, 'has_accelerate', False):
-                sd_model.to(devices.device)
-
+            else:
+                move_model(sd_model, devices.device)
             sd_models_compile.compile_diffusers(sd_model)
 
         if sd_model is None:
@@ -944,14 +953,14 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
         if hasattr(sd_model, "set_progress_bar_config"):
             sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba')
-        if op == 'refiner' and shared.opts.diffusers_move_refiner and not getattr(sd_model, 'has_accelerate', False):
+        if op == 'refiner' and shared.opts.diffusers_move_refiner:
             shared.log.debug('Moving refiner model to CPU')
-            sd_model.to(devices.cpu)
-        elif not getattr(sd_model, 'has_accelerate', False): # In offload modes, accelerate will move models around
-            sd_model.to(devices.device)
+            move_model(sd_model, devices.cpu)
+        else:
+            move_model(sd_model, devices.device)
         if op == 'refiner' and base_sent_to_cpu:
             shared.log.debug('Moving base model back to GPU')
-            model_data.sd_model.to(devices.device)
+            move_model(model_data.sd_model, devices.device)
     except Exception as e:
         shared.log.error("Failed to load diffusers model")
         errors.display(e, "loading Diffusers model")
@@ -1226,10 +1235,10 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
     else:
         shared.log.debug(f'Model weights loaded: {memory_stats()}')
     timer.record("load")
-    if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+    if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
         lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)
     else:
-        sd_model.to(devices.device)
+        move_model(sd_model, devices.device)
     timer.record("move")
     shared.log.debug(f'Model weights moved: {memory_stats()}')
     sd_hijack.model_hijack.hijack(sd_model)
@@ -1273,11 +1282,10 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
         current_checkpoint_info = getattr(sd_model, 'sd_checkpoint_info', None)
         if current_checkpoint_info is not None and checkpoint_info is not None and current_checkpoint_info.filename == checkpoint_info.filename:
             return None
-        if not getattr(sd_model, 'has_accelerate', False):
-            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
-                lowvram.send_everything_to_cpu()
-            else:
-                sd_model.to(devices.cpu)
+        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
+            lowvram.send_everything_to_cpu()
+        else:
+            move_model(sd_model, devices.cpu)
         if (reuse_dict or shared.opts.model_reuse_dict) and not getattr(sd_model, 'has_accelerate', False):
             shared.log.info('Reusing previous model dictionary')
             sd_hijack.model_hijack.undo_hijack(sd_model)
@@ -1322,8 +1330,8 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
         timer.record("hijack")
         script_callbacks.model_loaded_callback(sd_model)
         timer.record("callbacks")
-        if sd_model is not None and not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not getattr(sd_model, 'has_accelerate', False):
-            sd_model.to(devices.device)
+        if sd_model is not None and not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
+            move_model(sd_model, devices.device)
             timer.record("device")
     shared.state.end()
     shared.state = orig_state
@@ -1357,26 +1365,23 @@ def unload_model_weights(op='model'):
         if model_data.sd_model:
             if shared.backend == shared.Backend.ORIGINAL:
                 from modules import sd_hijack
-                model_data.sd_model.to(devices.cpu)
+                move_model(model_data.sd_model, devices.cpu)
                 sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
             elif not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
                 disable_offload(model_data.sd_model)
-                try:
-                    model_data.sd_model.to('meta')
-                except Exception:
-                    pass
+                move_model(model_data.sd_model, 'meta')
             model_data.sd_model = None
             devices.torch_gc(force=True)
             shared.log.debug(f'Unload weights {op}: {memory_stats()}')
-    else:
+    elif op == 'refiner':
         if model_data.sd_refiner:
             if shared.backend == shared.Backend.ORIGINAL:
                 from modules import sd_hijack
-                model_data.sd_model.to(devices.cpu)
+                move_model(model_data.sd_refiner, devices.cpu)
                 sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
             else:
-                disable_offload(model_data.sd_model)
-                model_data.sd_refiner.to('meta')
+                disable_offload(model_data.sd_refiner)
+                move_model(model_data.sd_refiner, 'meta')
             model_data.sd_refiner = None
             devices.torch_gc(force=True)
             shared.log.debug(f'Unload weights {op}: {memory_stats()}')

modules/sd_vae.py

@@ -239,11 +239,10 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         vae_source = "function-argument"
     if loaded_vae_file == vae_file:
         return None
-    if not getattr(sd_model, 'has_accelerate', False):
-        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
-            lowvram.send_everything_to_cpu()
-        else:
-            sd_model.to(devices.cpu)
+    if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
+        lowvram.send_everything_to_cpu()
+    else:
+        sd_models.move_model(sd_model, devices.cpu)
 
     if shared.backend == shared.Backend.ORIGINAL:
         sd_hijack.model_hijack.undo_hijack(sd_model)
@@ -260,6 +259,6 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
             if vae is not None:
                 sd_models.set_diffuser_options(sd_model, vae=vae, op='vae')
 
-    if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not getattr(sd_model, 'has_accelerate', False):
-        sd_model.to(devices.device)
+    if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
+        sd_models.move_model(sd_model, devices.device)
     return sd_model

modules/ui_interrogate.py

@@ -76,7 +76,7 @@ def interrogate_image(image, model, mode):
     shared.state.begin()
     shared.state.job = 'interrogate'
     try:
-        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
             lowvram.send_everything_to_cpu()
             devices.torch_gc()
         load_interrogator(model)
@@ -105,7 +105,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
     shared.state.job = 'batch interrogate'
     prompts = []
     try:
-        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
             lowvram.send_everything_to_cpu()
             devices.torch_gc()
         load_interrogator(model)

scripts/animatediff.py

@@ -97,8 +97,7 @@ def set_adapter(adapter_name: str = 'None'):
         )
         orig_pipe = shared.sd_model
         shared.sd_model = new_pipe
-        if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-            shared.sd_model.to(shared.device)
+        sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
         sd_models.copy_diffuser_options(new_pipe, orig_pipe)
         sd_models.set_diffuser_options(shared.sd_model, vae=None, op='model')
         shared.log.debug(f'AnimateDiff create pipeline: adapter="{loaded_adapter}"')

scripts/demofusion.py

@@ -14,7 +14,7 @@ from diffusers.schedulers import KarrasDiffusionSchedulers
 from diffusers.utils import is_accelerate_available, is_accelerate_version
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-from modules import scripts, processing, shared, sd_models
+from modules import scripts, processing, shared, sd_models, devices
 
 
 ### Class definition
@@ -1268,8 +1268,7 @@ class Script(scripts.Script):
             force_zeros_for_empty_prompt=shared.opts.diffusers_force_zeros,
         )
         shared.sd_model = new_pipe
-        if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-            shared.sd_model.to(shared.device)
+        sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
         sd_models.set_diffuser_options(shared.sd_model, vae=None, op='model')
         shared.log.debug(f'DemoFusion create: pipeline={shared.sd_model.__class__.__name__}')
         processed = processing.process_images(p)

scripts/example.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from diffusers.pipelines import StableDiffusionPipeline, StableDiffusionXLPipeline # pylint: disable=unused-import
-from modules import shared, scripts, processing, sd_models
+from modules import shared, scripts, processing, sd_models, devices
 
 """
 This is a simpler template for script for SD.Next that implements a custom pipeline
@@ -109,8 +109,8 @@ class Script(scripts.Script):
         )
         sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) # copy options from original pipeline
         sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc.
-        if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-            shared.sd_model.to(shared.device) # move pipeline if needed, but don't touch if its under automatic managment
+        sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
+        shared.sd_model.to(dtype=devices.dtype)
 
         # if pipeline also needs a specific type, you can set it here, but not commonly needed
         # shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)

scripts/image2video.py

@@ -85,8 +85,7 @@ class Script(scripts.Script):
             motion_adapter = diffusers.MotionAdapter.from_pretrained(repo_id)
             motion_adapter.to(devices.device, devices.dtype)
             shared.sd_model = sd_models.switch_pipe(diffusers.PIAPipeline, shared.sd_model, { 'motion_adapter': motion_adapter })
-            if not ((shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) or (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram)):
-                shared.sd_model.to(devices.device, devices.dtype)
+            sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
             if num_frames > 0:
                 p.task_args['num_frames'] = num_frames
                 p.task_args['image'] = p.init_images[0]
@@ -111,8 +110,8 @@ class Script(scripts.Script):
                 sd_models.copy_diffuser_options(pipe, shared.sd_model)
                 sd_models.set_diffuser_options(pipe)
                 shared.sd_model = pipe
-                shared.sd_model.to(devices.device, torch.float32)
-                devices.torch_gc()
+                sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
+                shared.sd_model.to(dtype=torch.float32)
             if num_frames > 0:
                 p.task_args['image'] = p.init_images[0]
                 p.task_args['num_frames'] = num_frames