code: def setup_for_low_vram(sd_model, use_medvram):
      if getattr(sd_model, "lowvram", False):
          return
  
      sd_model.lowvram = True
  
      parents = {}
  
      def send_me_to_gpu(module, _):
          """send this module to GPU; send whatever tracked module was previous in GPU to CPU;
          we add this as forward_pre_hook to a lot of modules and this way all but one of them will
          be in CPU
          """
          global module_in_gpu
  
          module = parents.get(module, module)
  
          if module_in_gpu == module:
              return
  
          if module_in_gpu is not None:
              module_in_gpu.to(cpu)
  
          module.to(devices.device)
          module_in_gpu = module
  
      # see below for register_forward_pre_hook;
      # first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is
      # useless here, and we just replace those methods
  
      first_stage_model = sd_model.first_stage_model
      first_stage_model_encode = sd_model.first_stage_model.encode
      first_stage_model_decode = sd_model.first_stage_model.decode
  
      def first_stage_model_encode_wrap(x):
          send_me_to_gpu(first_stage_model, None)
          return first_stage_model_encode(x)
  
      def first_stage_model_decode_wrap(z):
          send_me_to_gpu(first_stage_model, None)
          return first_stage_model_decode(z)
  
      to_remain_in_cpu = [
          (sd_model, "first_stage_model"),
          (sd_model, "depth_model"),
          (sd_model, "embedder"),
          (sd_model, "model"),
          (sd_model, "embedder"),
      ]
  
      is_sdxl = hasattr(sd_model, "conditioner")
      is_sd2 = not is_sdxl and hasattr(
          sd_model.cond_stage_model, "model"
      )
  
      if is_sdxl:
          to_remain_in_cpu.append((sd_model, "conditioner"))
      elif is_sd2:
          to_remain_in_cpu.append((sd_model.cond_stage_model, "model"))
      else:
          to_remain_in_cpu.append(
              (sd_model.cond_stage_model, "transformer")
          )
  
      # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model
      stored = []
      for obj, field in to_remain_in_cpu:
          module = getattr(obj, field, None)
          stored.append(module)
          setattr(obj, field, None)
  
      # send the model to GPU.
      sd_model.to(devices.device)
  
      # put modules back. the modules will be in CPU.
      for (obj, field), module in zip(to_remain_in_cpu, stored):
          setattr(obj, field, module)
  
      # register hooks for those the first three models
      if is_sdxl:
          sd_model.conditioner.register_forward_pre_hook(send_me_to_gpu)
      elif is_sd2:
          sd_model.cond_stage_model.model.register_forward_pre_hook(
              send_me_to_gpu
          )
          sd_model.cond_stage_model.model.token_embedding.register_forward_pre_hook(
              send_me_to_gpu
          )
          parents[
              sd_model.cond_stage_model.model
          ] = sd_model.cond_stage_model
          parents[
              sd_model.cond_stage_model.model.token_embedding
          ] = sd_model.cond_stage_model
      else:
          sd_model.cond_stage_model.transformer.register_forward_pre_hook(
              send_me_to_gpu
          )
          parents[
              sd_model.cond_stage_model.transformer
          ] = sd_model.cond_stage_model
  
      sd_model.first_stage_model.register_forward_pre_hook(
          send_me_to_gpu
      )
      sd_model.first_stage_model.encode = first_stage_model_encode_wrap
      sd_model.first_stage_model.decode = first_stage_model_decode_wrap
      if sd_model.depth_model:
          sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)
      if sd_model.embedder:
          sd_model.embedder.register_forward_pre_hook(send_me_to_gpu)
  
      if use_medvram:
          sd_model.model.register_forward_pre_hook(send_me_to_gpu)
      else:
          diff_model = sd_model.model.diffusion_model
  
          # the third remaining model is still too big for 4 GB, so we also do the same for its submodules
          # so that only one of them is in GPU at a time
          stored = (
              diff_model.input_blocks,
              diff_model.middle_block,
              diff_model.output_blocks,
              diff_model.time_embed,
          )
          (
              diff_model.input_blocks,
              diff_model.middle_block,
              diff_model.output_blocks,
              diff_model.time_embed,
          ) = (None, None, None, None)
          sd_model.model.to(devices.device)
          (
              diff_model.input_blocks,
              diff_model.middle_block,
              diff_model.output_blocks,
              diff_model.time_embed,
          ) = stored
  
          # install hooks for bits of third model
          diff_model.time_embed.register_forward_pre_hook(
              send_me_to_gpu
          )
          for block in diff_model.input_blocks:
              block.register_forward_pre_hook(send_me_to_gpu)
          diff_model.middle_block.register_forward_pre_hook(
              send_me_to_gpu
          )
          for block in diff_model.output_blocks:
              block.register_forward_pre_hook(send_me_to_gpu)
  
imports: import torch
  from modules import devices, shared
label_desc: Use image node_id 4 for calling a global function with example usage: setup_for_low_vram(sd_model, use_medvram) and returns: first_stage_model_encode, first_stage_model_decode