import gradio as gr import numpy as np from PIL import Image import io, os import urllib.request import torch import torch.nn as nn import torch.nn.functional as F # ── Download modelo ────────────────────────────────────────────────── MODEL_PATH = "RealESRGAN_x4plus.pth" MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth" if not os.path.exists(MODEL_PATH): print("Baixando modelo...") urllib.request.urlretrieve(MODEL_URL, MODEL_PATH) print("Pronto!") # ── Arquitetura RRDBNet (23 blocos — modelo correto) ───────────────── class ResidualDenseBlock(nn.Module): def __init__(self, num_feat=64, num_grow_ch=32): super().__init__() self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) self.conv2 = nn.Conv2d(num_feat+num_grow_ch, num_grow_ch, 3, 1, 1) self.conv3 = nn.Conv2d(num_feat+2*num_grow_ch, num_grow_ch, 3, 1, 1) self.conv4 = nn.Conv2d(num_feat+3*num_grow_ch, num_grow_ch, 3, 1, 1) self.conv5 = nn.Conv2d(num_feat+4*num_grow_ch, num_feat, 3, 1, 1) self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) def forward(self, x): x1=self.lrelu(self.conv1(x)) x2=self.lrelu(self.conv2(torch.cat((x,x1),1))) x3=self.lrelu(self.conv3(torch.cat((x,x1,x2),1))) x4=self.lrelu(self.conv4(torch.cat((x,x1,x2,x3),1))) return self.conv5(torch.cat((x,x1,x2,x3,x4),1))*0.2+x class RRDB(nn.Module): def __init__(self, num_feat=64, num_grow_ch=32): super().__init__() self.rdb1=ResidualDenseBlock(num_feat,num_grow_ch) self.rdb2=ResidualDenseBlock(num_feat,num_grow_ch) self.rdb3=ResidualDenseBlock(num_feat,num_grow_ch) def forward(self, x): out=self.rdb1(x); out=self.rdb2(out); out=self.rdb3(out) return out*0.2+x class RRDBNet(nn.Module): def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32): super().__init__() self.conv_first=nn.Conv2d(num_in_ch,num_feat,3,1,1) self.body=nn.Sequential(*[RRDB(num_feat,num_grow_ch) for _ in range(num_block)]) self.conv_body=nn.Conv2d(num_feat,num_feat,3,1,1) self.conv_up1=nn.Conv2d(num_feat,num_feat,3,1,1) self.conv_up2=nn.Conv2d(num_feat,num_feat,3,1,1) self.conv_hr=nn.Conv2d(num_feat,num_feat,3,1,1) self.conv_last=nn.Conv2d(num_feat,num_out_ch,3,1,1) self.lrelu=nn.LeakyReLU(negative_slope=0.2,inplace=True) def forward(self, x): feat=self.conv_first(x) feat=feat+self.conv_body(self.body(feat)) feat=self.lrelu(self.conv_up1(F.interpolate(feat,scale_factor=2,mode='nearest'))) feat=self.lrelu(self.conv_up2(F.interpolate(feat,scale_factor=2,mode='nearest'))) return self.conv_last(self.lrelu(self.conv_hr(feat))) # ── Carrega pesos com strict=True (correto para 23 blocos) ─────────── device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Usando: {device}") model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32) state=torch.load(MODEL_PATH, map_location=device, weights_only=False) if 'params_ema' in state: weights=state['params_ema'] elif 'params' in state: weights=state['params'] else: weights=state model.load_state_dict(weights, strict=True) model.eval() model = model.to(device) if device.type == 'cuda': model = model.half() # float16 na GPU — 2x mais rápido torch.set_num_threads(4) print("Modelo carregado!") # ── Tile adaptativo — tile menor para imagens grandes ──────────────── def get_tile_size(w, h): if device.type == 'cuda': return 512, 16 # GPU tem VRAM suficiente para tiles grandes px = w * h if px < 200000: return 128, 8 if px < 500000: return 96, 6 return 64, 5 def infer_tile(img_t): b,c,h,w = img_t.shape TILE, PAD = get_tile_size(w, h) img_t = img_t.to(device) if device.type == 'cuda': img_t = img_t.half() if h <= TILE and w <= TILE: with torch.inference_mode(): out = model(img_t) return out.float().cpu() output = torch.zeros(b, c, h*4, w*4) y = 0 while y < h: x = 0 while x < w: y0=max(y-PAD,0); x0=max(x-PAD,0) y1=min(y+TILE+PAD,h); x1=min(x+TILE+PAD,w) with torch.inference_mode(): tile_out=model(img_t[:,:,y0:y1,x0:x1]) tile_out = tile_out.float().cpu() pt=(y-y0)*4; pl=(x-x0)*4 pb=(y1-min(y+TILE,h))*4; pr=(x1-min(x+TILE,w))*4 oy0=y*4; ox0=x*4 oy1=min(y+TILE,h)*4; ox1=min(x+TILE,w)*4 sy1=tile_out.shape[2]-(pb if pb>0 else 0) sx1=tile_out.shape[3]-(pr if pr>0 else 0) output[:,:,oy0:oy1,ox0:ox1]=tile_out[:,:,pt:sy1,pl:sx1] x+=TILE y+=TILE return output # ── Função principal ───────────────────────────────────────────────── def upscale_image(input_image): try: if input_image is None: raise gr.Error("Nenhuma imagem enviada") pil_img = input_image.convert('RGB') w, h = pil_img.size print(f"Processando {w}×{h}...") MAX_PX = 1500 if w > MAX_PX or h > MAX_PX: raise gr.Error(f"Imagem muito grande ({w}×{h}). Máximo: {MAX_PX}×{MAX_PX} px.") img_np = np.array(pil_img).astype(np.float32)/255.0 img_t = torch.from_numpy(img_np).permute(2,0,1).unsqueeze(0) out_t = infer_tile(img_t) out_np = (out_t.squeeze(0).permute(1,2,0).clamp(0,1).numpy()*255).astype(np.uint8) result = Image.fromarray(out_np) print(f"Resultado: {result.width}×{result.height}") # Salva PNG sem compressão em arquivo temporário out_path = f"/tmp/result_{os.getpid()}.png" result.save(out_path, format='PNG', compress_level=3) return out_path except gr.Error: raise except Exception as e: raise gr.Error(str(e)) # ── Interface Gradio ───────────────────────────────────────────────── with gr.Blocks(title="Iluminados Upscaler") as demo: gr.Markdown("## Real-ESRGAN ×4 — Iluminados Upscaler API") with gr.Row(): inp = gr.Image(label="Entrada", type="pil") out = gr.File(label="Resultado ×4 (PNG)") gr.Button("Upscale").click(fn=upscale_image, inputs=inp, outputs=out) demo.launch()