Update app.py

app.py

@@ -9,36 +9,51 @@ import tempfile
 import requests
 
 # ——————————————————————————————————————————————————————————————
-# 1) Path ke ONNX model “×2” dan “×4” (diunduh dari GitHub RAW)
+# 1) Path ke ONNX model “×2” dan “×4” di folder model/
 MODEL_DIR     = "model"
 MODEL_X2_PATH = os.path.join(MODEL_DIR, "Real-ESRGAN_x2plus.onnx")
 MODEL_X4_PATH = os.path.join(MODEL_DIR, "Real-ESRGAN-x4plus.onnx")
 
-def download_model(path: str, url: str):
+# File ID Google Drive (contoh—ganti dengan ID asli Anda)
+FILE_ID_X2 = "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6"
+FILE_ID_X4 = "1wDBHad9RCJgJDGsPdapLYl3cr8j-PMJ6"
+
+def download_from_drive(file_id: str, dest_path: str):
     """
-    Jika file model belum ada, unduh dari URL RAW GitHub.
+    Download file besar dari Google Drive dengan mengatasi prompt konfirmasi (jika ada).
     """
-    if not os.path.isfile(path):
-        os.makedirs(MODEL_DIR, exist_ok=True)
-        response = requests.get(url, stream=True)
-        if response.status_code == 200:
-            with open(path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    if chunk:
-                        f.write(chunk)
-            print(f"Model telah diunduh dan disimpan di {path}")
-        else:
-            raise FileNotFoundError(f"Model gagal diunduh dari {url}, status code: {response.status_code}")
-
-# **Ganti link “blob” dengan raw.githubusercontent.com**  
-download_model(
-    MODEL_X2_PATH,
-    "https://raw.githubusercontent.com/salmanalfarisi11/Upscaler_images/master/model/Real-ESRGAN_x2plus.onnx"
-)
-download_model(
-    MODEL_X4_PATH,
-    "https://raw.githubusercontent.com/salmanalfarisi11/Upscaler_images/master/model/Real-ESRGAN-x4plus.onnx"
-)
+    URL = "https://drive.google.com/uc?export=download"
+    session = requests.Session()
+
+    # Permintaan awal
+    response = session.get(URL, params={"id": file_id}, stream=True)
+    token = None
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            token = value
+            break
+
+    # Jika membutuhkan konfirmasi, ulangi dengan token
+    if token:
+        params = {"id": file_id, "confirm": token}
+        response = session.get(URL, params=params, stream=True)
+
+    # Tuliskan file ke disk
+    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+    with open(dest_path, "wb") as f:
+        for chunk in response.iter_content(chunk_size=32768):
+            if chunk:
+                f.write(chunk)
+    print(f"Model telah diunduh dan disimpan di {dest_path}")
+    return dest_path
+
+# Unduh model ×2
+if not os.path.isfile(MODEL_X2_PATH):
+    download_from_drive(FILE_ID_X2, MODEL_X2_PATH)
+
+# Unduh model ×4
+if not os.path.isfile(MODEL_X4_PATH):
+    download_from_drive(FILE_ID_X4, MODEL_X4_PATH)
 
 # ——————————————————————————————————————————————————————————————
 # 2) Buat ONNXRuntime session hanya dengan CPU, 2 thread
@@ -74,7 +89,7 @@ if SCALE_X4 != 4:
     raise RuntimeError(f"Model ×4 menghasilkan scale = {SCALE_X4}, bukan 4")
 
 # ——————————————————————————————————————————————————————————————
-# 3) Fungsi util untuk mem‐proses satu tile untuk ×2 dan ×4
+# 3) Fungsi util untuk mem‐proses satu tile
 def run_tile_x2(tile_np: np.ndarray) -> np.ndarray:
     patch_nchw = np.transpose(tile_np, (2, 0, 1))[None, ...]
     out_nchw = session_x2.run(None, {input_meta_x2.name: patch_nchw})[0]
@@ -97,6 +112,7 @@ def tile_upscale(input_img: Image.Image, scale: int, max_dim=1024):
     else:
         H_in, W_in, run_tile, SCALE = H_in_x4, W_in_x4, run_tile_x4, SCALE_X4
 
+    # Resize jika terlalu besar
     w, h = input_img.size
     if w > max_dim or h > max_dim:
         scale_factor = max_dim / float(max(w, h))
@@ -168,7 +184,7 @@ css = """
 """
 
 # ——————————————————————————————————————————————————————————————
-# 7) Bangun Gradio Blocks – tiga tombol (×2, ×4, ×8)
+# 7) Bangun Gradio Blocks
 with gr.Blocks(css=css, title="Real-ESRGAN Triple-Mode Upscaler") as demo:
     gr.Markdown(
         """