Update app.py

app.py

@@ -1,71 +1,174 @@
-import gradio as gr
 import os
 import zipfile
+import tempfile
+from pathlib import Path
+
 import cv2
-import urllib.request
+import gradio as gr
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
 
-# Baixa o modelo específico para estética anime/ilustração se ele não estiver no servidor
-model_path = 'RealESRGAN_x4plus_anime_6B.pth'
-if not os.path.isfile(model_path):
-    url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth'
-    urllib.request.urlretrieve(url, model_path)
+# =========================
+# CONFIG
+# =========================
+OUTSCALE = 2  # saída final em 2x
+
+HF_REPO_ID = os.getenv("HF_REPO_ID", "xinntao/Real-ESRGAN")
+HF_FILENAME = os.getenv("HF_FILENAME", "RealESRGAN_x4plus_anime_6B.pth")
+CACHE_DIR = os.getenv("HF_HOME", "/tmp/hf-cache")
+
+
+def download_model() -> str:
+    return hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=HF_FILENAME,
+        cache_dir=CACHE_DIR,
+    )
+
+
+def build_upsampler():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    use_half = device == "cuda"
+
+    model_path = download_model()
+
+    # Modelo anime 6B
+    model = RRDBNet(
+        num_in_ch=3,
+        num_out_ch=3,
+        num_feat=64,
+        num_block=6,
+        num_grow_ch=32,
+        scale=4,
+    )
+
+    upsampler = RealESRGANer(
+        scale=4,
+        model_path=model_path,
+        model=model,
+        tile=256,
+        tile_pad=10,
+        pre_pad=0,
+        half=use_half,
+        device=device,
+    )
+    return upsampler
+
+
+UPSAMPLER = build_upsampler()
+
+
+def upscale_one_image(image: np.ndarray) -> np.ndarray:
+    if image is None:
+        return None
+
+    if image.dtype != np.uint8:
+        image = np.clip(image, 0, 1)
+        image = (image * 255).astype(np.uint8)
+
+    # Suporte a alpha
+    if image.ndim == 3 and image.shape[2] == 4:
+        rgb = image[:, :, :3]
+        alpha = image[:, :, 3]
+
+        bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
+        out_bgr, _ = UPSAMPLER.enhance(bgr, outscale=OUTSCALE)
+        out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+
+        alpha_up = cv2.resize(
+            alpha,
+            (alpha.shape[1] * OUTSCALE, alpha.shape[0] * OUTSCALE),
+            interpolation=cv2.INTER_CUBIC,
+        )
+        return np.dstack([out_rgb, alpha_up])
 
-# Configura a arquitetura neural
-model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+    bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    out_bgr, _ = UPSAMPLER.enhance(bgr, outscale=OUTSCALE)
+    out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+    return out_rgb
 
-# Inicializa o motor de Upscale
-upsampler = RealESRGANer(
-    scale=4,
-    model_path=model_path,
-    model=model,
-    half=False # Essencial manter como False para rodar no plano grátis (CPU)
-)
 
 def process_batch(files):
+    """
+    Recebe uma lista de arquivos, processa um por um e retorna:
+    - galeria com previews
+    - caminho do zip final
+    """
     if not files:
-        return None
-        
-    out_dir = "imagens_prontas"
-    os.makedirs(out_dir, exist_ok=True)
-    
-    # Limpa a pasta antes de uma nova rodada para não misturar os lotes
-    for f in os.listdir(out_dir):
-        os.remove(os.path.join(out_dir, f))
-
-    # Processa cada imagem enviada
-    for file in files:
-        filename = os.path.basename(file.name)
-        img = cv2.imread(file.name, cv2.IMREAD_UNCHANGED)
-        
-        if img is not None:
-            # Aplica o upscale 4x
-            output, _ = upsampler.enhance(img, outscale=4)
-            cv2.imwrite(os.path.join(out_dir, filename), output)
-
-    # Cria o arquivo ZIP final
-    zip_filename = "Imagens_Alta_Resolucao.zip"
-    with zipfile.ZipFile(zip_filename, 'w') as zipf:
-        for folderName, subfolders, filenames in os.walk(out_dir):
-            for filename in filenames:
-                filePath = os.path.join(folderName, filename)
-                zipf.write(filePath, os.path.basename(filePath))
-                
-    return zip_filename
-
-# Constrói a interface visual do celular
-with gr.Blocks(theme=gr.themes.Soft()) as app:
-    gr.Markdown("# 🚀 Upscale em Lote")
-    gr.Markdown("Envie várias imagens de uma vez. O servidor processará o upscale preservando a arte original e entregará um arquivo ZIP.")
-    
+        return [], None
+
+    previews = []
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+        out_dir = tmpdir / "upscaled"
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        for file_obj in files:
+            in_path = Path(file_obj.name)
+            image = cv2.imread(str(in_path), cv2.IMREAD_UNCHANGED)
+
+            if image is None:
+                continue
+
+            # OpenCV lê em BGR/BGRA; converter para RGB/RGBA para o pipeline
+            if image.ndim == 3 and image.shape[2] == 4:
+                image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
+            elif image.ndim == 3:
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+            result = upscale_one_image(image)
+            if result is None:
+                continue
+
+            out_name = f"{in_path.stem}_2x.png"
+            out_path = out_dir / out_name
+
+            if result.ndim == 3 and result.shape[2] == 4:
+                # RGBA -> BGRA para salvar via OpenCV
+                save_img = cv2.cvtColor(result, cv2.COLOR_RGBA2BGRA)
+            else:
+                # RGB -> BGR para salvar via OpenCV
+                save_img = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+
+            cv2.imwrite(str(out_path), save_img)
+
+            previews.append((result, out_name))
+
+        zip_path = tmpdir / "upscaled_images.zip"
+        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+            for img_file in out_dir.iterdir():
+                zf.write(img_file, arcname=img_file.name)
+
+        # Copia o zip para um caminho persistente temporário do Gradio
+        final_zip = Path(tempfile.gettempdir()) / "upscaled_images.zip"
+        final_zip.write_bytes(zip_path.read_bytes())
+
+        return previews, str(final_zip)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Anime Upscaler 2x\nUpload em lote com saída em ZIP.")
+
     with gr.Row():
-        with gr.Column():
-            image_inputs = gr.File(label="1. Selecione as imagens", file_count="multiple", type="filepath")
-            btn_process = gr.Button("2. Iniciar Upscale", variant="primary")
-        with gr.Column():
-            zip_output = gr.File(label="3. Download do ZIP Pronto")
-            
-    btn_process.click(fn=process_batch, inputs=image_inputs, outputs=zip_output)
-
-app.launch()
+        files_in = gr.Files(
+            label="Envie várias imagens",
+            file_types=["image"],
+            file_count="multiple",
+        )
+
+    run_btn = gr.Button("Processar")
+    gallery_out = gr.Gallery(label="Prévia", columns=2, height=420)
+    zip_out = gr.File(label="Baixar ZIP")
+
+    run_btn.click(
+        fn=process_batch,
+        inputs=files_in,
+        outputs=[gallery_out, zip_out],
+    )
+
+if __name__ == "__main__":
+    demo.launch()