Merge pull request #5595 from wywywywy/ldsr-safetensors (1b66d527) · Commits · github_fork / Stable Diffusion Webui

extensions-builtin/LDSR/ldsr_model_arch.py

+8 −2

Original line number	Diff line number	Diff line
		import os
		import gc
		import time
		import warnings
		@@ -8,6 +9,7 @@ import torchvision
		from PIL import Image
		from einops import rearrange, repeat
		from omegaconf import OmegaConf
		import safetensors.torch

		from ldm.models.diffusion.ddim import DDIMSampler
		from ldm.util import instantiate_from_config, ismap
		@@ -28,8 +30,12 @@ class LDSR:
		model: torch.nn.Module = cached_ldsr_model
		else:
		print(f"Loading model from {self.modelPath}")
		_, extension = os.path.splitext(self.modelPath)
		if extension.lower() == ".safetensors":
		pl_sd = safetensors.torch.load_file(self.modelPath, device="cpu")
		else:
		pl_sd = torch.load(self.modelPath, map_location="cpu")
		sd = pl_sd["state_dict"]
		sd = pl_sd["state_dict"] if "state_dict" in pl_sd else pl_sd
		config = OmegaConf.load(self.yamlPath)
		config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1"
		model: torch.nn.Module = instantiate_from_config(config.model)

+6 −2

Original line number	Diff line number	Diff line
		@@ -25,6 +25,7 @@ class UpscalerLDSR(Upscaler):
		yaml_path = os.path.join(self.model_path, "project.yaml")
		old_model_path = os.path.join(self.model_path, "model.pth")
		new_model_path = os.path.join(self.model_path, "model.ckpt")
		safetensors_model_path = os.path.join(self.model_path, "model.safetensors")
		if os.path.exists(yaml_path):
		statinfo = os.stat(yaml_path)
		if statinfo.st_size >= 10485760:
		@@ -33,6 +34,9 @@ class UpscalerLDSR(Upscaler):
		if os.path.exists(old_model_path):
		print("Renaming model from model.pth to model.ckpt")
		os.rename(old_model_path, new_model_path)
		if os.path.exists(safetensors_model_path):
		model = safetensors_model_path
		else:
		model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
		file_name="model.ckpt", progress=True)
		yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,