Spaces:
Running on Zero
Running on Zero
| """Stage I one-step diffusion model (Sec. 3.1): Stable Diffusion 2.1 with a pixel LoRA and a detail LoRA, adapted from | |
| PiSA-SR (third_party/PiSA-SR). The restored latent is z_hat = z_blur - eps(z_blur), with eps the U-Net prediction at | |
| timestep 1 and an empty prompt, and | |
| eps = eps_pix + omega * (eps_pix+det - eps_pix) | |
| where eps_pix uses the pixel LoRA alone, eps_pix+det both LoRAs and omega is the detail weight. omega = 1 is a single | |
| U-Net pass with both LoRAs merged; any other omega needs both predictions (two passes). | |
| With separate_unets=True the two predictions come from two U-Nets (pixel; pixel + detail) instead of one U-Net whose | |
| weights are swapped in place between the passes: twice the U-Net memory, the same output, and every model is an | |
| ordinary module, as needed to move it between devices (e.g. Hugging Face ZeroGPU) or to wrap it. | |
| PiSA-SR names the detail LoRA "sem" (semantic) and omega "lambda_sem"; these names remain in its checkpoint keys.""" | |
| import argparse | |
| import copy | |
| import torch | |
| from .. import paths | |
| from .._compat import patch_huggingface_hub | |
| patch_huggingface_hub() # before diffusers is imported | |
| NO_TILING = 100000 # tile sizes above any input: the whole image is processed at once | |
| def _pisa_sr_eval_class(): | |
| paths.use_third_party(paths.PISA_SR) | |
| from pisasr import PiSASR_eval | |
| return PiSASR_eval | |
| class OneStepDiffusion(_pisa_sr_eval_class()): | |
| """PiSA-SR's evaluation model with the changes needed for deblurring: | |
| * deterministic VAE encoding (latent_dist.mode()), so the restored image is reproducible; | |
| * the Stage I extras of our checkpoints (trained U-Net conv_in, LoRA on the VAE encoder); | |
| * the merged pixel + detail weights of the two-pass mode are built once and reused.""" | |
| def _load_pretrained_weights(self, pretrained_path): | |
| sd = torch.load(pretrained_path, map_location="cpu") | |
| if "conv_in" in sd: # base weights, before the pixel LoRA is merged on top | |
| self.unet.load_state_dict(sd["conv_in"], strict=False) | |
| self._load_and_save_ckpt_from_state_dict(sd) | |
| if "state_dict_vae" in sd: | |
| from peft import LoraConfig | |
| self.vae.add_adapter(LoraConfig(r=sd["vae_lora_rank"], init_lora_weights="gaussian", | |
| target_modules=sd["vae_lora_encoder_modules"]), adapter_name="vae_enc") | |
| for n, p in self.vae.named_parameters(): | |
| if n in sd["state_dict_vae"]: | |
| p.data.copy_(sd["state_dict_vae"][n]) | |
| def _apply_lora_delta(self): | |
| # back to the pixel + detail weights after the pixel-only pass, without re-adding the deltas every call | |
| if getattr(self, "_merged_weight", None) is None: | |
| self._merged_weight = {n: self.lora_deltas_sem[n] + w if n in self.lora_deltas_sem else w | |
| for n, w in self.ori_unet_weight.items()} | |
| for name, param in self.unet.named_parameters(): | |
| param.data = self._merged_weight[name] | |
| def detail_weight(self): | |
| """omega""" | |
| return float(self.lambda_sem) | |
| def detail_weight(self, omega): | |
| assert not self.single_pass or omega == 1.0, "load the model with detail_weight != 1 to change it" | |
| self.lambda_sem = torch.tensor([float(omega)], device=self.device) | |
| def uses_single_pass(self): | |
| """One U-Net pass when loaded with omega = 1, or when omega is set back to 1.""" | |
| return self.single_pass or float(self.lambda_sem) == 1.0 | |
| def empty_prompt_embedding(self): | |
| return self.encode_prompt([""]).to(dtype=self.weight_dtype) | |
| def predict_noise(self, z, prompt_embeds): | |
| """eps for the latent z (whole latent, no tiling).""" | |
| if not getattr(self, "separate_unets", False): | |
| return self._predict_no_tiling(z, prompt_embeds, self.uses_single_pass) | |
| t = torch.ones(1, dtype=torch.long, device=z.device) # timestep 1 | |
| eps = self.unet(z, t, encoder_hidden_states=prompt_embeds).sample | |
| if self.uses_single_pass: | |
| return eps | |
| eps_pixel = self.unet_pixel(z, t, encoder_hidden_states=prompt_embeds).sample | |
| eps -= eps_pixel # same operations as PiSA-SR's two passes | |
| return self.lambda_pix.to(z.device) * eps_pixel + self.lambda_sem.to(z.device) * eps | |
| def use_separate_unets(self): | |
| """Second U-Net with the pixel LoRA only (unet keeps pixel + detail); drops the weight tables of the swap.""" | |
| self.unet_pixel = copy.deepcopy(self.unet) | |
| for name, p in self.unet_pixel.named_parameters(): | |
| p.data = self.ori_unet_weight[name].clone() | |
| del self.ori_unet_weight, self.lora_deltas_sem, self._merged_weight | |
| self.separate_unets = True | |
| def load_one_step_diffusion(checkpoint=paths.STAGE1, sd21=paths.SD21, detail_weight=1.0, precision="fp16", | |
| merge_vae_lora=False, separate_unets=False): | |
| """Stage I in eval mode. detail_weight != 1 selects the two-pass mode, in which it can be changed later; 1 merges | |
| the detail LoRA into a single pass. | |
| merge_vae_lora folds the VAE-encoder LoRA into the encoder weights: the encoder runs ~20% faster, but the output is | |
| no longer bit-identical to the reported results (about 60 dB PSNR apart), so it is off by default.""" | |
| single = detail_weight == 1.0 | |
| args = argparse.Namespace(pretrained_model_path=str(sd21), pretrained_path=str(checkpoint), lambda_pix=1.0, | |
| lambda_sem=detail_weight, mixed_precision=precision, default=single, | |
| vae_decoder_tiled_size=NO_TILING, vae_encoder_tiled_size=NO_TILING, | |
| latent_tiled_size=NO_TILING, latent_tiled_overlap=0, seed=42) | |
| model = OneStepDiffusion(args) | |
| model.set_eval() | |
| model.single_pass = single | |
| if not single: | |
| # Use the pixel + detail weights that every two-pass call restores after its pixel-only pass from the start. | |
| # peft's own merge differs from them in the last bits, which made the first image restored differ slightly | |
| # (up to 2/255) from later ones. | |
| model._apply_lora_delta() | |
| if separate_unets: | |
| model.use_separate_unets() | |
| if merge_vae_lora: | |
| from peft.tuners.lora import LoraLayer | |
| for mod in model.vae.modules(): | |
| if isinstance(mod, LoraLayer): | |
| mod.merge() | |
| return model | |
| def stage1_parameter_count(model, refiner=None): | |
| """Parameters of everything a configuration runs: text encoder, U-Net with merged LoRAs, VAE, and the refiner | |
| without the VAE encoder it does not run.""" | |
| n = sum(p.numel() for m in (model.text_encoder, model.unet, model.vae) for p in m.parameters()) | |
| if refiner is not None: | |
| n += sum(v.numel() for k, v in refiner.state_dict().items() if not k.startswith(("vae.encoder", "vae.quant_conv"))) | |
| return n | |