Spaces:
Running
on
Zero
Running
on
Zero
Update utils.py
Browse files
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 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
|
156 |
|
157 |
-
def upscale(upscale_model, tensor: torch.Tensor, inf_device, output_device="cpu") -> torch.Tensor:
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
|
173 |
-
|
174 |
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
|
185 |
-
|
186 |
-
|
187 |
|
188 |
|
189 |
-
def upscale_batch_and_concatenate(upscale_model, latents, inf_device, output_device="cpu") -> torch.Tensor:
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
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"])
|