tsqn commited on
Commit
2801bd7
·
verified ·
1 Parent(s): 8e39b51

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +45 -39
utils.py CHANGED
@@ -8,10 +8,12 @@ import numpy as np
8
  import itertools
9
  import PIL.Image
10
  import safetensors.torch
 
 
11
  import tqdm
12
  import logging
13
  from diffusers.utils import export_to_video
14
- from spandrel import ModelLoader
15
 
16
  logger = logging.getLogger(__file__)
17
 
@@ -146,53 +148,53 @@ def tiled_scale(
146
  )
147
 
148
 
149
- def load_sd_upscale(ckpt, inf_device):
150
- sd = load_torch_file(ckpt, device=inf_device)
151
- if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
152
- sd = state_dict_prefix_replace(sd, {"module.": ""})
153
- out = ModelLoader().load_from_state_dict(sd).half()
154
- return out
155
 
156
 
157
- def upscale(upscale_model, tensor: torch.Tensor, inf_device, output_device="cpu") -> torch.Tensor:
158
- memory_required = module_size(upscale_model.model)
159
- memory_required += (
160
- (512 * 512 * 3) * tensor.element_size() * max(upscale_model.scale, 1.0) * 384.0
161
- ) # The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
162
- memory_required += tensor.nelement() * tensor.element_size()
163
- print(f"UPScaleMemory required: {memory_required / 1024 / 1024 / 1024} GB")
164
 
165
- upscale_model.to(inf_device)
166
- tile = 512
167
- overlap = 32
168
 
169
- steps = tensor.shape[0] * get_tiled_scale_steps(
170
- tensor.shape[3], tensor.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
171
- )
172
 
173
- pbar = ProgressBar(steps, desc="Tiling and Upscaling")
174
 
175
- s = tiled_scale(
176
- samples=tensor.to(torch.float16),
177
- function=lambda a: upscale_model(a),
178
- tile_x=tile,
179
- tile_y=tile,
180
- overlap=overlap,
181
- upscale_amount=upscale_model.scale,
182
- pbar=pbar,
183
- )
184
 
185
- upscale_model.to(output_device)
186
- return s
187
 
188
 
189
- def upscale_batch_and_concatenate(upscale_model, latents, inf_device, output_device="cpu") -> torch.Tensor:
190
- upscaled_latents = []
191
- for i in range(latents.size(0)):
192
- latent = latents[i]
193
- upscaled_latent = upscale(upscale_model, latent, inf_device, output_device)
194
- upscaled_latents.append(upscaled_latent)
195
- return torch.stack(upscaled_latents)
196
 
197
 
198
  def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8):
@@ -219,3 +221,7 @@ class ProgressBar:
219
 
220
  # 更新进度
221
  self.b_unit.update(self.current)
 
 
 
 
 
8
  import itertools
9
  import PIL.Image
10
  import safetensors.torch
11
+ import subprocess
12
+ import sys
13
  import tqdm
14
  import logging
15
  from diffusers.utils import export_to_video
16
+ #from spandrel import ModelLoader
17
 
18
  logger = logging.getLogger(__file__)
19
 
 
148
  )
149
 
150
 
151
+ # def load_sd_upscale(ckpt, inf_device):
152
+ # sd = load_torch_file(ckpt, device=inf_device)
153
+ # if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
154
+ # sd = state_dict_prefix_replace(sd, {"module.": ""})
155
+ # out = ModelLoader().load_from_state_dict(sd).half()
156
+ # return out
157
 
158
 
159
+ # def upscale(upscale_model, tensor: torch.Tensor, inf_device, output_device="cpu") -> torch.Tensor:
160
+ # memory_required = module_size(upscale_model.model)
161
+ # memory_required += (
162
+ # (512 * 512 * 3) * tensor.element_size() * max(upscale_model.scale, 1.0) * 384.0
163
+ # ) # The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
164
+ # memory_required += tensor.nelement() * tensor.element_size()
165
+ # print(f"UPScaleMemory required: {memory_required / 1024 / 1024 / 1024} GB")
166
 
167
+ # upscale_model.to(inf_device)
168
+ # tile = 512
169
+ # overlap = 32
170
 
171
+ # steps = tensor.shape[0] * get_tiled_scale_steps(
172
+ # tensor.shape[3], tensor.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
173
+ # )
174
 
175
+ # pbar = ProgressBar(steps, desc="Tiling and Upscaling")
176
 
177
+ # s = tiled_scale(
178
+ # samples=tensor.to(torch.float16),
179
+ # function=lambda a: upscale_model(a),
180
+ # tile_x=tile,
181
+ # tile_y=tile,
182
+ # overlap=overlap,
183
+ # upscale_amount=upscale_model.scale,
184
+ # pbar=pbar,
185
+ # )
186
 
187
+ # upscale_model.to(output_device)
188
+ # return s
189
 
190
 
191
+ # def upscale_batch_and_concatenate(upscale_model, latents, inf_device, output_device="cpu") -> torch.Tensor:
192
+ # upscaled_latents = []
193
+ # for i in range(latents.size(0)):
194
+ # latent = latents[i]
195
+ # upscaled_latent = upscale(upscale_model, latent, inf_device, output_device)
196
+ # upscaled_latents.append(upscaled_latent)
197
+ # return torch.stack(upscaled_latents)
198
 
199
 
200
  def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8):
 
221
 
222
  # 更新进度
223
  self.b_unit.update(self.current)
224
+
225
+ def install_packages():
226
+ subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "pip", "wheel", "setuptools"])
227
+ subprocess.check_call([sys.executable, "-m", "pip", "install", "-pre", "torch", "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/nightly/cu124"])