Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,759 Bytes
b273838 |
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 |
import gradio as gr
import numpy as np
import spaces
import torch
import torchvision
from latent_utils import compress
from util.file import generate_binary_file, load_numpy_from_binary_bitwise
from util.img_utils import resize_and_crop
@torch.no_grad()
@spaces.GPU(duration=80)
def main(img_to_compress, T, K, model_type='512x512', bitstream=None, avail_models=None,
progress=gr.Progress(track_tqdm=True)):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
indices = load_numpy_from_binary_bitwise(bitstream, K, T, model_type, T - 1)
if indices is not None:
indices = indices.to(device)
if indices is None:
img_to_compress = resize_and_crop(img_to_compress, int(model_type.split('x')[0]))
img_to_compress = (torchvision.transforms.ToTensor()(img_to_compress) * 2) - 1
img_to_compress = img_to_compress.unsqueeze(0).to(device)
else:
img_to_compress = None
print(T, K, model_type)
# model, _ = load_model(img_size_to_id[img_size], T, device, float16=True, compile=False)
model = avail_models[model_type].to(device)
model.device = device
model.model.to(device=device)
model.model.scheduler.device = device
# model.model.scheduler.scheduler = model.model.scheduler.scheduler.to(device)
model.set_timesteps(T, device=device)
model.num_timesteps = T
with torch.no_grad():
x, indices = compress(model, img_to_compress, K, indices, device=device)
x = (x / 2 + 0.5).clamp(0, 1)
x = x.detach().cpu().squeeze().numpy()
x = np.transpose(x, (1, 2, 0))
torch.cuda.empty_cache()
indices = generate_binary_file(indices.numpy(), K, T, model_type)
if bitstream is None:
return x, indices
return x
|