Update app.py

app.py

@@ -11,6 +11,8 @@ import time
 import types
 import cv2
 import urllib.request
+import shutil
+import tempfile
 import warnings
 import torch
 import numpy as np
@@ -95,28 +97,51 @@ print("Setting up GFPGAN & Real-ESRGAN...")
 weights_dir = "weights"
 os.makedirs(weights_dir, exist_ok=True)
 
+# Download into a temporary file and atomically publish only a complete weight file.
+def ensure_weight_file(path, url, min_bytes=1_000_000):
+    if os.path.isfile(path) and os.path.getsize(path) >= min_bytes:
+        return
+
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    fd, temp_path = tempfile.mkstemp(prefix=".download-", dir=os.path.dirname(path) or ".")
+    os.close(fd)
+    try:
+        with urllib.request.urlopen(url, timeout=60) as response, open(temp_path, "wb") as out:
+            shutil.copyfileobj(response, out)
+        if os.path.getsize(temp_path) < min_bytes:
+            raise RuntimeError(f"Downloaded weight file is unexpectedly small: {path}")
+        os.replace(temp_path, path)
+    finally:
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+
+
 # 1. Real-ESRGAN Weights Setup
 esrgan_path = f"{weights_dir}/RealESRGAN_x4plus.pth"
-if not os.path.exists(esrgan_path):
-    url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
-    urllib.request.urlretrieve(url, esrgan_path)
+ensure_weight_file(
+    esrgan_path,
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+)
 
 esrgan_model = RRDBNet(
     num_in_ch=3, num_out_ch=3, num_feat=64,
     num_block=23, num_grow_ch=32, scale=4
 )
 
+restoration_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 enhancer = RealESRGANer(
     scale=4, model_path=esrgan_path, model=esrgan_model,
-    tile=512, tile_pad=10, pre_pad=0, half=True,
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    tile=512, tile_pad=10, pre_pad=0,
+    half=(restoration_device.type == "cuda"),
+    device=restoration_device,
 )
 
 # 2. GFPGAN Weights Setup
 gfpgan_path = f"{weights_dir}/GFPGANv1.3.pth"
-if not os.path.exists(gfpgan_path):
-    url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
-    urllib.request.urlretrieve(url, gfpgan_path)
+ensure_weight_file(
+    gfpgan_path,
+    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
+)
 
 gfpgan_restorer = GFPGANer(
     model_path=gfpgan_path,
@@ -124,7 +149,7 @@ gfpgan_restorer = GFPGANer(
     arch="clean",
     channel_multiplier=2,
     bg_upsampler=enhancer,
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device=restoration_device
 )
 
 print("GFPGAN & Real-ESRGAN ready.")
@@ -158,10 +183,21 @@ print("All models loaded successfully.")
 # 6. HELPER FUNCTIONS
 # ================================================================
 def is_greyscale_or_sepia(image):
-    rgb = np.array(image.convert("RGB"))
+    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
     hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
-    sat = np.mean(hsv[:, :, 1])
-    return (sat < 25), False
+    mean_sat = float(np.mean(hsv[:, :, 1]))
+    is_greyscale = mean_sat < 25
+
+    # A simple warm-channel heuristic; intentionally avoid classifying grayscale
+    # images as sepia. This is a prompt hint, not a strict image classifier.
+    channel_means = rgb.mean(axis=(0, 1))
+    red_mean, green_mean, blue_mean = channel_means
+    is_sepia = (
+        not is_greyscale
+        and red_mean > green_mean + 8
+        and green_mean > blue_mean + 5
+    )
+    return is_greyscale, bool(is_sepia)
 
 
 def calculate_padding(ow, oh, direction, pct):
@@ -174,7 +210,9 @@ def calculate_padding(ow, oh, direction, pct):
     elif direction == "Bottom": bottom = v_ext
     elif direction == "Horizontal": left = right = h_ext
     elif direction == "Vertical": top = bottom = v_ext
-    elif direction in ["Both", "All Sides"]: left = right = top = bottom = h_ext
+    elif direction in ["Both", "All Sides"]:
+        left = right = h_ext
+        top = bottom = v_ext
     return left, top, right, bottom
 
 
@@ -340,6 +378,9 @@ def fast_single_pass_outpaint(
         raise gr.Error("Please upload an image first.")
     
     start_time = time.time()
+    steps = int(steps)
+    if not 10 <= steps <= 50:
+        raise gr.Error("Inference steps must be between 10 and 50.")
 
     image = cap_image(image, 1024)
     ow, oh = image.size
@@ -370,6 +411,20 @@ def fast_single_pass_outpaint(
     target_w = int(16 * math.ceil(target_w / 16))
     target_h = int(16 * math.ceil(target_h / 16))
 
+    # Bound worst-case FLUX memory use. Raise this only after validating the
+    # selected ZeroGPU allocation can handle larger canvases reliably.
+    max_outpaint_side = 2048
+    max_outpaint_pixels = 4_194_304
+    if (
+        target_w > max_outpaint_side
+        or target_h > max_outpaint_side
+        or target_w * target_h > max_outpaint_pixels
+    ):
+        raise gr.Error(
+            f"Requested canvas is too large ({target_w}x{target_h}). "
+            f"Reduce the extension percentage or use a smaller input image."
+        )
+
     canvas = Image.new("RGB", (target_w, target_h), (127, 127, 127))
     canvas.paste(image, (left, top))
 
@@ -429,15 +484,19 @@ def fast_single_pass_outpaint(
                     bs = brisque(t, data_range=1.0).item()
                     final_info += f"\n📉 BRISQUE: {bs:.2f} (lower better)"
 
-                clip_model.to("cuda")
-                with torch.no_grad():
-                    inputs = clip_processor(
-                        text=[prompt[:77]], images=current_img,
-                        return_tensors="pt", padding=True, truncation=True
-                    ).to("cuda")
-                    clip_score = clip_model(**inputs).logits_per_image.item()
-                clip_model.to("cpu")
-                final_info += f"\n🤖 CLIP: {clip_score:.4f} (higher better)"
+                try:
+                    clip_model.to("cuda")
+                    with torch.no_grad():
+                        inputs = clip_processor(
+                            text=[prompt[:77]], images=current_img,
+                            return_tensors="pt", padding=True, truncation=True
+                        ).to("cuda")
+                        clip_score = clip_model(**inputs).logits_per_image.item()
+                    final_info += f"\n🤖 CLIP: {clip_score:.4f} (higher better)"
+                finally:
+                    # Keep this large optional metrics model off the GPU even
+                    # when preprocessing or scoring fails.
+                    clip_model.to("cpu")
             except Exception as e:
                 final_info += f"\n⚠️ Metrics error: {e}"