shared_init.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import os
  2. import torch
  3. from modules import shared
  4. from modules.shared import cmd_opts
  5. def initialize():
  6. """Initializes fields inside the shared module in a controlled manner.
  7. Should be called early because some other modules you can import mingt need these fields to be already set.
  8. """
  9. os.makedirs(cmd_opts.hypernetwork_dir, exist_ok=True)
  10. from modules import options, shared_options
  11. shared.options_templates = shared_options.options_templates
  12. shared.opts = options.Options(shared_options.options_templates, shared_options.restricted_opts)
  13. shared.restricted_opts = shared_options.restricted_opts
  14. try:
  15. shared.opts.load(shared.config_filename)
  16. except FileNotFoundError:
  17. pass
  18. from modules import devices
  19. devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = \
  20. (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])
  21. devices.dtype = torch.float32 if cmd_opts.no_half else torch.float16
  22. devices.dtype_vae = torch.float32 if cmd_opts.no_half or cmd_opts.no_half_vae else torch.float16
  23. devices.dtype_inference = torch.float32 if cmd_opts.precision == 'full' else devices.dtype
  24. if cmd_opts.precision == "half":
  25. msg = "--no-half and --no-half-vae conflict with --precision half"
  26. assert devices.dtype == torch.float16, msg
  27. assert devices.dtype_vae == torch.float16, msg
  28. assert devices.dtype_inference == torch.float16, msg
  29. devices.force_fp16 = True
  30. devices.force_model_fp16()
  31. shared.device = devices.device
  32. shared.weight_load_location = None if cmd_opts.lowram else "cpu"
  33. from modules import shared_state
  34. shared.state = shared_state.State()
  35. from modules import styles
  36. shared.prompt_styles = styles.StyleDatabase(shared.styles_filename)
  37. from modules import interrogate
  38. shared.interrogator = interrogate.InterrogateModels("interrogate")
  39. from modules import shared_total_tqdm
  40. shared.total_tqdm = shared_total_tqdm.TotalTQDM()
  41. from modules import memmon, devices
  42. shared.mem_mon = memmon.MemUsageMonitor("MemMon", devices.device, shared.opts)
  43. shared.mem_mon.start()