|
|
import spaces |
|
|
import numpy as np |
|
|
import gradio as gr |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from models.arch.RDnet_ import FullNet_NLP |
|
|
from models.arch.classifier import PretrainedConvNext |
|
|
import torchvision.transforms.functional as TF |
|
|
|
|
|
class Pipe: |
|
|
def __init__(self): |
|
|
channels = [64, 128, 256, 512] |
|
|
layers = [2, 2, 4, 2] |
|
|
num_subnet = 4 |
|
|
self.net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3) |
|
|
for param in self.net_i.parameters(): |
|
|
param.data = param.data.to(torch.float16) |
|
|
self.net_i.load_state_dict(torch.load('./fp16_check.pt')['icnn']) |
|
|
self.net_i = self.net_i.to('cpu') |
|
|
self.net_c = PretrainedConvNext("convnext_small_in22k") |
|
|
self.net_c.load_state_dict(torch.load('./classifier_32.pt')['icnn']) |
|
|
self.net_c=self.net_c.to('cpu') |
|
|
|
|
|
self.net_i.eval().to('cuda') |
|
|
self.net_c.eval().to('cuda') |
|
|
self.output = None |
|
|
|
|
|
def __call__(self, img): |
|
|
with torch.no_grad(): |
|
|
image_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) |
|
|
h, w = image_tensor.shape[-2], image_tensor.shape[-1] |
|
|
h, w = h // 32 * 32, w // 32 * 32 |
|
|
image_tensor = torch.nn.functional.interpolate(image_tensor, size=(h, w), mode='bilinear').cuda() |
|
|
|
|
|
ipt=self.net_c(image_tensor) |
|
|
image_tensor = image_tensor.half() |
|
|
ipt = ipt.half() |
|
|
output_i, output_j=self.net_i(image_tensor,ipt,prompt=True) |
|
|
clean = output_j[-1][:, 3:, ...] |
|
|
clean=torch.clamp(clean, 0, 1) |
|
|
|
|
|
self.output = clean |
|
|
|
|
|
pipe = Pipe() |
|
|
|
|
|
|
|
|
@spaces.GPU(duration=120) |
|
|
def predict(img): |
|
|
pipe(img) |
|
|
return pipe.output |
|
|
demo=gr.Interface(predict, gr.Image(), "image") |
|
|
|
|
|
demo.launch() |
|
|
|
