Spaces:
Running on Zero
Running on Zero
File size: 4,377 Bytes
5c331a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Tiled per-component SR with measured CUDA costs and bounded concurrent waves."""
import gc
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from PIL import Image
from .resources import parallelism, snapshot
def upscale_parts(parts, model, log):
import torch
tile = 256
device = "cuda"
def predict(image):
array = np.array(image.convert("RGB"), dtype=np.float32) / 255
tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).to(device)
with torch.inference_mode():
output = model(tensor).clamp(0, 1)
return Image.fromarray((output[0].permute(1, 2, 0).cpu().numpy() * 255).round().astype(np.uint8))
# Profile one maximal padded tile before admitting simultaneous requests.
while True:
try:
torch.cuda.empty_cache()
free, _ = torch.cuda.mem_get_info()
before = torch.cuda.memory_allocated()
torch.cuda.reset_peak_memory_stats()
predict(Image.new("RGB", (tile + 32, tile + 32)))
torch.cuda.synchronize()
per_gpu = max(torch.cuda.max_memory_allocated() - before, 64 * 1024**2)
per_ram = max(max(p.visual.material.baseColorTexture.width *
p.visual.material.baseColorTexture.height * 4 * 16 * 4 for _, p in parts),
128 * 1024**2)
parallelism(snapshot(free), per_ram, per_gpu, len(parts))
break
except (torch.cuda.OutOfMemoryError, MemoryError):
gc.collect()
torch.cuda.empty_cache()
if tile <= 32:
raise MemoryError("A single minimum-size SR tile cannot fit with the 10% reserve.") from None
tile //= 2
def worker(item):
key, part = item
source = part.visual.material.baseColorTexture.convert("RGBA")
output = Image.new("RGBA", (source.width * 4, source.height * 4))
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
for y in range(0, source.height, tile):
for x in range(0, source.width, tile):
right, bottom = min(x + tile, source.width), min(y + tile, source.height)
left_pad, top_pad = max(0, x - 16), max(0, y - 16)
right_pad, bottom_pad = min(source.width, right + 16), min(source.height, bottom + 16)
patch = predict(source.crop((left_pad, top_pad, right_pad, bottom_pad)))
patch = patch.crop(((x-left_pad)*4, (y-top_pad)*4,
(right-left_pad)*4, (bottom-top_pad)*4))
output.paste(patch, (x * 4, y * 4))
stream.synchronize()
output.putalpha(source.getchannel("A").resize(output.size, Image.Resampling.BILINEAR))
return key, output
remaining = list(parts)
completed = {}
cap = len(parts)
while remaining:
torch.cuda.empty_cache()
free, _ = torch.cuda.mem_get_info()
count = min(cap, parallelism(snapshot(free), per_ram, per_gpu, len(remaining)))
log(f"Upscaling {count} component(s) concurrently; tile={tile}, VRAM/task={per_gpu}, reserve=10%")
wave, remaining = remaining[:count], remaining[count:]
retries = []
with ThreadPoolExecutor(max_workers=count) as pool:
futures = [(item, pool.submit(worker, item)) for item in wave]
for item, future in futures:
try:
key, texture = future.result()
completed[key] = texture
except torch.cuda.OutOfMemoryError as error:
error.__traceback__ = None
retries.append(item)
if retries:
gc.collect()
torch.cuda.empty_cache()
if count == 1:
if tile <= 32:
raise MemoryError("Texture SR ran out of memory at minimum tile size.")
tile //= 2
cap = max(1, count // 2)
remaining = retries + remaining
for key, part in parts:
part.visual.material.baseColorTexture = completed[key]
return {"tile_size": tile, "measured_vram_per_task": per_gpu, "ram_budget_per_task": per_ram,
"reserve": 0.10, "scale": 4, "model": "RealESRGAN_x4plus", "component_count": len(parts)}
|