Update app.py

app.py

@@ -2,40 +2,42 @@ import os
 import zipfile
 import tempfile
 from pathlib import Path
+from urllib.request import urlretrieve
 
 import cv2
 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
 
 # =========================
 # CONFIG
 # =========================
-OUTSCALE = 2  # saída final em 2x
+OUTSCALE = 2
 
-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")
+# Peso oficial do anime model mostrado no README do Real-ESRGAN
+MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+MODEL_DIR = Path("weights")
+MODEL_PATH = MODEL_DIR / "RealESRGAN_x4plus_anime_6B.pth"
 
 
-def download_model() -> str:
-    return hf_hub_download(
-        repo_id=HF_REPO_ID,
-        filename=HF_FILENAME,
-        cache_dir=CACHE_DIR,
-    )
+def ensure_model() -> str:
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+    if MODEL_PATH.exists() and MODEL_PATH.stat().st_size > 0:
+        return str(MODEL_PATH)
+
+    urlretrieve(MODEL_URL, MODEL_PATH)
+    return str(MODEL_PATH)
 
 
 def build_upsampler():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     use_half = device == "cuda"
 
-    model_path = download_model()
+    model_path = ensure_model()
 
-    # Modelo anime 6B
     model = RRDBNet(
         num_in_ch=3,
         num_out_ch=3,
@@ -45,7 +47,7 @@ def build_upsampler():
         scale=4,
     )
 
-    upsampler = RealESRGANer(
+    return RealESRGANer(
         scale=4,
         model_path=model_path,
         model=model,
@@ -55,7 +57,6 @@ def build_upsampler():
         half=use_half,
         device=device,
     )
-    return upsampler
 
 
 UPSAMPLER = build_upsampler()
@@ -69,7 +70,6 @@ def upscale_one_image(image: np.ndarray) -> np.ndarray:
         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]
@@ -92,11 +92,6 @@ def upscale_one_image(image: np.ndarray) -> np.ndarray:
 
 
 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
 
@@ -114,7 +109,6 @@ def process_batch(files):
             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:
@@ -128,14 +122,11 @@ def process_batch(files):
             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"
@@ -143,7 +134,6 @@ def process_batch(files):
             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())
 
@@ -151,14 +141,13 @@ def process_batch(files):
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Anime Upscaler 2x\nUpload em lote com saída em ZIP.")
+    gr.Markdown("# Anime Upscaler 2x\nUpload em lote e baixe um ZIP.")
 
-    with gr.Row():
-        files_in = gr.Files(
-            label="Envie várias imagens",
-            file_types=["image"],
-            file_count="multiple",
-        )
+    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)