update

mikazuki/utils/devices.py

@@ -20,6 +20,7 @@ def check_torch_gpu():
  printable_devices.append(f"GPU {pos}: {torch.cuda.get_device_name(device)} ({round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024 / 1024)} GB)")
  log.info(f'Torch detected GPU: {torch.cuda.get_device_name(device)} VRAM {round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024)} Arch {torch.cuda.get_device_capability(device)} Cores {torch.cuda.get_device_properties(device).multi_processor_count}')
  else:
- log.warn("Torch is not able to use GPU, please check your torch installation.\n Use --skip-prepare-environment to disable this check")
+ log.error("Torch is not able to use GPU, please check your torch installation.\n Use --skip-prepare-environment to disable this check")
+ log.error("！！！Torch 无法使用 GPU，您无法正常开始训练！！！\n您的显卡可能并不支持，或是 torch 安装有误。请检查您的 torch 安装。\n 使用 --skip-prepare-environment 可以跳过此检查")
  except Exception as e:
  log.error(f'Could not load torch: {e}')

mikazuki/utils/train_utils.py

@@ -18,6 +18,18 @@ def is_promopt_like(s):
 
 def validate_model(model_name: str):
  if os.path.exists(model_name):
+ try:
+ with open(model_name, "rb") as f:
+ content = f.read(1024 * 200)
+ if b"model.diffusion_model" in content or b"cond_stage_model.transformer.text_model" in content:
+ return True, "ok"
+
+ if b"lora_unet" in content or b"lora_te" in content:
+ return False, "pretrained model is not a Stable Diffusion checkpoint / 校验失败：底模不是 Stable Diffusion 模型"
+ except Exception as e:
+ log.warn(f"model file {model_name} can't open: {e}")
+ return True, ""
+
  return True, "ok"
 
  # huggerface model repo