Update app.py

app.py

@@ -3,7 +3,7 @@
 import os
 import sys
 
-# --- Install Dependencies (ensure diffusers, gradio_imageslider, huggingface-hub are present) ---
+# --- Install Dependencies ---
 print("Installing required packages: diffusers, gradio_imageslider, huggingface-hub…")
 os.system("pip install --no-input diffusers gradio_imageslider huggingface-hub")
 
@@ -22,7 +22,7 @@ from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
 from gradio_imageslider import ImageSlider
 from PIL import Image, ImageOps
-from huggingface_hub import snapshot_download, HFValidationError
+from huggingface_hub import snapshot_download
 
 # --- Logging & Device Setup ---
 logging.basicConfig(level=logging.INFO)
@@ -68,14 +68,16 @@ try:
         repo_id=flux_model_id,
         repo_type="model",
         local_dir=local_model_dir,
-        use_auth_token=huggingface_token,       # ← pass token here
+        use_auth_token=huggingface_token,
         ignore_patterns=["*.md", "*.gitattributes"],
     )
     logging.info(f"Downloaded base model to: {model_path}")
 
     logging.info(f"Loading ControlNet: {controlnet_model_id}")
     controlnet = FluxControlNetModel.from_pretrained(
-        controlnet_model_id, torch_dtype=torch_dtype, use_auth_token=huggingface_token
+        controlnet_model_id,
+        torch_dtype=torch_dtype,
+        use_auth_token=huggingface_token
     ).to(device)
     logging.info("ControlNet loaded.")
 
@@ -88,14 +90,13 @@ try:
     ).to(device)
     logging.info("Pipeline ready.")
 
-except HFValidationError as hf_err:
-    logging.error(f"Gated‑model access error: {hf_err}", exc_info=True)
-    print(f"FATAL: Cannot access gated repo {flux_model_id}. "
-          "Have you accepted its license and set a valid HUGGINGFACE_TOKEN?")
-    sys.exit(1)
 except Exception as e:
     logging.error(f"Error loading models: {e}", exc_info=True)
-    print(f"FATAL: model load failed: {e}")
+    # Detect gated‐repo 401 by checking for “401” in the message
+    if "401" in str(e):
+        print(f"FATAL: Cannot access gated model {flux_model_id}. Have you accepted its license and set a valid HUGGINGFACE_TOKEN?")
+    else:
+        print(f"FATAL: model load failed: {e}")
     sys.exit(1)
 
 # --- Constants & Helpers ---
@@ -110,7 +111,8 @@ def process_input(input_image):
     if img.mode != "RGB":
         img = img.convert("RGB")
     w, h = img.size
-    # enforce intermediate budget
+
+    # enforce intermediate‐scale budget
     target_px = (w*INTERNAL_PROCESSING_FACTOR)*(h*INTERNAL_PROCESSING_FACTOR)
     if target_px > MAX_PIXEL_BUDGET:
         max_in = MAX_PIXEL_BUDGET / (INTERNAL_PROCESSING_FACTOR**2)
@@ -120,11 +122,13 @@ def process_input(input_image):
         was_resized = True
     else:
         was_resized = False
-    # round to multiple of 8
+
+    # round dimensions to multiples of 8
     w2, h2 = img.size
     w2 -= w2 % 8; h2 -= h2 % 8
     if img.size != (w2,h2):
         img = img.resize((w2,h2), Image.Resampling.LANCZOS)
+
     return img, w, h, was_resized
 
 @spaces.GPU(duration=75)
@@ -149,7 +153,7 @@ def infer(
     processed, w0, h0, resized_flag = process_input(input_image)
     w_proc, h_proc = processed.size
 
-    # make control image
+    # prepare control image at INTERNAL scale
     cw, ch = w_proc*INTERNAL_PROCESSING_FACTOR, h_proc*INTERNAL_PROCESSING_FACTOR
     control_img = processed.resize((cw, ch), Image.Resampling.LANCZOS)
 
@@ -165,7 +169,7 @@ def infer(
             generator=gen
         ).images[0]
 
-    # final resize
+    # final resize to user factor
     if resized_flag:
         fw, fh = w_proc*final_upscale_factor, h_proc*final_upscale_factor
     else:
@@ -189,17 +193,17 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
         with gr.Column(scale=2):
             inp = gr.Image(label="Input Image", type="pil", sources=["upload","clipboard"], height=350)
         with gr.Column(scale=1):
-            upf = gr.Slider("Final Upscale Factor", 1, INTERNAL_PROCESSING_FACTOR, step=1, value=2)
+            upf   = gr.Slider("Final Upscale Factor", 1, INTERNAL_PROCESSING_FACTOR, step=1, value=2)
             steps = gr.Slider("Inference Steps", 4, 50, step=1, value=15)
-            cscale = gr.Slider("ControlNet Scale", 0.0, 1.5, step=0.05, value=0.6)
+            cscale= gr.Slider("ControlNet Scale", 0.0, 1.5, step=0.05, value=0.6)
             with gr.Row():
                 sld = gr.Slider("Seed", 0, MAX_SEED, step=1, value=42)
                 rnd = gr.Checkbox("Randomize", value=True, scale=0, min_width=80)
             btn = gr.Button("⚡ Upscale Image", variant="primary")
 
     slider = ImageSlider("Input / Output", type="pil", interactive=False, show_label=True, position=0.5)
-    out_seed = gr.Textbox("Seed Used", interactive=False, visible=True)
-    out_b64  = gr.Textbox("API Base64 Output", interactive=False, visible=False)
+    out_seed= gr.Textbox("Seed Used", interactive=False, visible=True)
+    out_b64 = gr.Textbox("API Base64 Output", interactive=False, visible=False)
 
     btn.click(
         fn=infer,
@@ -208,5 +212,5 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
         api_name="upscale"
     )
 
-# expose JSON API at /run/upscale
+# Expose JSON API at /run/upscale
 demo.queue(max_size=10).launch(share=False, show_api=True)