app.py

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()