Update app.py

app.py

@@ -1,4 +1,3 @@
-# ---- Imports ----
 import logging
 import random
 import warnings
@@ -7,9 +6,9 @@ import io
 import base64
 import gradio as gr
 import numpy as np
-import spaces
+import spaces # Ensure spaces is imported for the decorator
 import torch
-# --- Add this if you want to try Solution 2 later ---
+# --- Optional: Add this if trying Solution 2 (suppress_errors) later ---
 # import torch._dynamo
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
@@ -18,7 +17,6 @@ from PIL import Image, ImageOps
 from huggingface_hub import snapshot_download
 
 # --- Setup Logging and Device ---
-# ... (rest of setup code remains the same) ...
 logging.basicConfig(level=logging.INFO)
 warnings.filterwarnings("ignore")
 
@@ -28,9 +26,13 @@ css = """
 """
 
 if torch.cuda.is_available():
-    power_device = "GPU"; device = "cuda"; torch_dtype = torch.bfloat16
+    power_device = "GPU"
+    device = "cuda"
+    torch_dtype = torch.bfloat16
 else:
-    power_device = "CPU"; device = "cpu"; torch_dtype = torch.float32
+    power_device = "CPU"
+    device = "cpu"
+    torch_dtype = torch.float32
 logging.info(f"Selected device: {device} | Data type: {torch_dtype}")
 
 # --- Authentication and Model Download ---
@@ -41,25 +43,36 @@ local_model_dir = flux_model_id.split('/')[-1]
 pipe = None
 
 try:
-    # ... (model download code remains the same) ...
     logging.info(f"Downloading base model: {flux_model_id}")
-    model_path = snapshot_download(repo_id=flux_model_id, repo_type="model", ignore_patterns=["*.md", "*.gitattributes"], local_dir=local_model_dir, token=huggingface_token)
+    model_path = snapshot_download(
+        repo_id=flux_model_id,
+        repo_type="model",
+        ignore_patterns=["*.md", "*.gitattributes"],
+        local_dir=local_model_dir,
+        token=huggingface_token
+    )
     logging.info(f"Base model downloaded/verified in: {model_path}")
 
     logging.info(f"Loading ControlNet model: {controlnet_model_id}")
-    controlnet = FluxControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch_dtype).to(device)
+    controlnet = FluxControlNetModel.from_pretrained(
+        controlnet_model_id, torch_dtype=torch_dtype
+    ).to(device)
     logging.info("ControlNet model loaded.")
 
     logging.info("Loading FluxControlNetPipeline...")
-    pipe = FluxControlNetPipeline.from_pretrained(model_path, controlnet=controlnet, torch_dtype=torch_dtype)
+    pipe = FluxControlNetPipeline.from_pretrained(
+        model_path,
+        controlnet=controlnet,
+        torch_dtype=torch_dtype
+    )
     pipe.to(device)
     logging.info("Pipeline loaded and moved to device.")
 
     # --- OPTIMIZATION: Attempt torch.compile (PyTorch 2.0+) ---
     if device == "cuda" and hasattr(torch, "compile"):
-        # --- TRY THIS FIRST: Change mode to 'default' ---
+        # --- Using "default" mode as first attempt to fix compile errors ---
         compile_mode = "default"
-        # --- Alternative (Solution 2): Uncomment these lines ---
+        # --- Alternative (Solution 2): Uncomment these lines if "default" fails ---
         # import torch._dynamo
         # torch._dynamo.config.suppress_errors = True
         # compile_mode = "max-autotune" # or "default" even with suppress_errors
@@ -69,261 +82,403 @@ try:
         try:
             pipe.transformer = torch.compile(pipe.transformer, mode=compile_mode, fullgraph=True)
             logging.info("Pipeline transformer compiled successfully.")
-            # Optional dummy inference run can go here
+            # Optional dummy inference run can go here for pre-compilation
         except Exception as e:
             logging.warning(f"torch.compile failed (mode='{compile_mode}'): {e}. Running unoptimized.")
-            # --- Solution 3: If compilation fails consistently, comment out the compile line above ---
+            # --- Fallback (Solution 3): If compilation fails consistently, comment out the compile line above ---
             # pipe.transformer = torch.compile(pipe.transformer, mode=compile_mode, fullgraph=True) # <-- Comment this out
     else:
         logging.info("torch.compile not available or not on CUDA, skipping compilation.")
 
-    # --- (Optional xformers code would go here if used) ---
+    # --- (Optional xformers code could be added here if used) ---
 
     logging.info("Pipeline ready for inference.")
 
 except Exception as e:
+    # Log the full error traceback for debugging
     logging.error(f"FATAL: Error during model loading or setup: {e}", exc_info=True)
+    # Print a simple message to console as well
     print(f"FATAL ERROR DURING MODEL LOAD/SETUP: {e}")
+    # Exit if models can't load - Gradio app won't work anyway
     raise SystemExit(f"Model loading/setup failed: {e}")
 
-
 # --- Constants ---
 MAX_SEED = 2**32 - 1
-MAX_PIXEL_BUDGET = 1280 * 1280
-INTERNAL_PROCESSING_FACTOR = 4
+MAX_PIXEL_BUDGET = 1280 * 1280 # Max pixels for the *intermediate* image
+INTERNAL_PROCESSING_FACTOR = 4 # Factor used for quality generation
 
 # --- Image Processing Function (process_input) ---
-# ... (process_input function remains the same) ...
 def process_input(input_image):
-    if input_image is None: raise gr.Error("Input image is missing!")
+    """Processes input image: handles orientation, converts to RGB, checks budget, rounds dimensions."""
+    if input_image is None:
+        raise gr.Error("Input image is missing!")
     try:
+        # Ensure it's a PIL Image and handle EXIF orientation
         input_image = ImageOps.exif_transpose(input_image)
-        if input_image.mode != 'RGB': input_image = input_image.convert('RGB')
+        # Convert to RGB if needed
+        if input_image.mode != 'RGB':
+            logging.info(f"Converting input image from {input_image.mode} to RGB")
+            input_image = input_image.convert('RGB')
         w, h = input_image.size
-    except AttributeError: raise gr.Error("Invalid input image format.")
-    except Exception as img_err: raise gr.Error(f"Could not process input image: {img_err}")
+    except AttributeError:
+        # Catch cases where input_image might not be a valid PIL Image object
+        raise gr.Error("Invalid input image format. Please provide a valid image file.")
+    except Exception as img_err:
+         # Catch other potential PIL errors
+         raise gr.Error(f"Could not process input image: {img_err}")
 
     w_original, h_original = w, h
-    if w == 0 or h == 0: raise gr.Error("Input image has zero dimensions.")
+    if w == 0 or h == 0:
+        raise gr.Error("Input image has zero width or height.")
 
+    # Calculate target intermediate size based on INTERNAL factor for budget check
     target_w_internal = w * INTERNAL_PROCESSING_FACTOR
     target_h_internal = h * INTERNAL_PROCESSING_FACTOR
     target_pixels_internal = target_w_internal * target_h_internal
+
     was_resized = False
-    input_image_to_process = input_image.copy()
+    input_image_to_process = input_image.copy() # Work on a copy
 
+    # Check if the *intermediate* size exceeds the budget
     if target_pixels_internal > MAX_PIXEL_BUDGET:
+        # Calculate the maximum allowed input pixels for the given internal factor
         max_input_pixels = MAX_PIXEL_BUDGET / (INTERNAL_PROCESSING_FACTOR**2)
         current_input_pixels = w * h
+
         if current_input_pixels > max_input_pixels:
+            # Calculate scaling factor to fit budget and resize input
             input_scale_factor = (max_input_pixels / current_input_pixels) ** 0.5
-            input_w_resized = max(8, int(w * input_scale_factor))
-            input_h_resized = max(8, int(h * input_scale_factor))
+            input_w_resized = max(8, int(w * input_scale_factor)) # Ensure min size 8
+            input_h_resized = max(8, int(h * input_scale_factor)) # Ensure min size 8
             intermediate_w = input_w_resized * INTERNAL_PROCESSING_FACTOR
             intermediate_h = input_h_resized * INTERNAL_PROCESSING_FACTOR
-            logging.warning(f"Requested {INTERNAL_PROCESSING_FACTOR}x intermediate exceeds budget. Resizing input {w}x{h} -> {input_w_resized}x{input_h_resized}.")
-            gr.Info(f"Intermediate {INTERNAL_PROCESSING_FACTOR}x size exceeds budget. Input resized to {input_w_resized}x{input_h_resized} -> model generates ~{int(intermediate_w)}x{int(intermediate_h)}.")
+
+            logging.warning(
+                f"Requested {INTERNAL_PROCESSING_FACTOR}x intermediate output ({target_w_internal}x{target_h_internal}) exceeds budget. "
+                f"Resizing input from {w}x{h} to {input_w_resized}x{input_h_resized}."
+            )
+            gr.Info(
+                f"Intermediate {INTERNAL_PROCESSING_FACTOR}x size exceeds budget. Input resized to {input_w_resized}x{input_h_resized} "
+                f"-> model generates ~{int(intermediate_w)}x{int(intermediate_h)}."
+            )
             input_image_to_process = input_image_to_process.resize((input_w_resized, input_h_resized), Image.Resampling.LANCZOS)
-            was_resized = True
+            was_resized = True # Flag that original dimensions were lost for final scaling
 
+    # Round processed input dimensions down to nearest multiple of 8 (required by some models)
     w_proc, h_proc = input_image_to_process.size
-    w_final_proc = max(8, w_proc - w_proc % 8)
-    h_final_proc = max(8, h_proc - h_proc % 8)
+    w_final_proc = max(8, w_proc - w_proc % 8) # Ensure minimum 8x8
+    h_final_proc = max(8, h_proc - h_proc % 8) # Ensure minimum 8x8
+
     if (w_proc, h_proc) != (w_final_proc, h_final_proc):
-        logging.info(f"Rounding processed input dims {w_proc}x{h_proc} -> {w_final_proc}x{h_final_proc}")
+        logging.info(f"Rounding processed input dimensions from {w_proc}x{h_proc} to {w_final_proc}x{h_final_proc}")
         input_image_to_process = input_image_to_process.resize((w_final_proc, h_final_proc), Image.Resampling.LANCZOS)
 
     return input_image_to_process, w_original, h_original, was_resized
 
 
 # --- Inference Function (infer) ---
-@spaces.GPU(duration=180)
+# --- MODIFIED GPU DURATION ---
+@spaces.GPU(duration=75)
 def infer(
-    seed, randomize_seed, input_image, num_inference_steps,
-    final_upscale_factor, controlnet_conditioning_scale,
-    progress=gr.Progress(track_tqdm=True),
+    seed,
+    randomize_seed,
+    input_image,
+    num_inference_steps,
+    final_upscale_factor,
+    controlnet_conditioning_scale,
+    progress=gr.Progress(track_tqdm=True), # Gradio progress tracking
 ):
+    """Runs the Flux ControlNet upscaling pipeline."""
     global pipe
-    # --- IMPROVED ERROR HANDLING: Define default return early ---
-    default_return = [[input_image, None], int(seed) if seed is not None else 0, None]
+    # --- Define default return structure for error cases ---
+    # [[Input Image or None, Output Image or None], Seed Int, Base64 String or None]
+    current_seed = int(seed) if seed is not None else 0
+    default_return = [[input_image, None], current_seed, None]
 
     if pipe is None:
         gr.Error("Pipeline not loaded. Cannot perform inference.")
-        return default_return # Use default return
+        return default_return
 
-    original_input_pil = input_image # Keep ref even if None initially
+    original_input_pil = input_image # Keep reference to original input
 
+    # Handle missing input image
     if input_image is None:
          gr.Warning("Please provide an input image.")
-         # Update seed in default return if randomized
-         if randomize_seed: seed = random.randint(0, MAX_SEED)
-         else: seed = int(seed) if seed is not None else 0
-         default_return[1] = seed
-         return default_return # Use default return
-
-    if randomize_seed: seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-    # --- UPDATE DEFAULT RETURN SEED ---
+         # Update seed in default return if randomized, keep original image as None
+         if randomize_seed: current_seed = random.randint(0, MAX_SEED)
+         default_return[1] = current_seed
+         default_return[0][0] = None # Explicitly set original image part to None
+         return default_return
+
+    # Determine seed
+    if randomize_seed:
+        current_seed = random.randint(0, MAX_SEED)
+    seed = int(current_seed) # Use the final determined seed
+    # Update default return with final seed and original image
     default_return[1] = seed
-    # --- Ensure original image is in the default return ---
     default_return[0][0] = original_input_pil
 
-
+    # Ensure numerical inputs are integers
     final_upscale_factor = int(final_upscale_factor)
     num_inference_steps = int(num_inference_steps)
 
+    # Clamp final factor if needed
     if final_upscale_factor > INTERNAL_PROCESSING_FACTOR:
         gr.Warning(f"Clamping final upscale factor {final_upscale_factor}x to internal {INTERNAL_PROCESSING_FACTOR}x.")
         final_upscale_factor = INTERNAL_PROCESSING_FACTOR
 
-    logging.info(f"Starting inference: seed={seed}, internal={INTERNAL_PROCESSING_FACTOR}x, final={final_upscale_factor}x, steps={num_inference_steps}, cnet_scale={controlnet_conditioning_scale}")
+    logging.info(
+        f"Starting inference: seed={seed}, internal={INTERNAL_PROCESSING_FACTOR}x, "
+        f"final={final_upscale_factor}x, steps={num_inference_steps}, "
+        f"cnet_scale={controlnet_conditioning_scale}"
+    )
 
+    # Process the input image
     try:
-        processed_input_image, w_original, h_original, was_input_resized = process_input(input_image)
+        processed_input_image, w_original, h_original, was_input_resized = process_input(
+            input_image
+        )
     except Exception as e:
         logging.error(f"Error processing input image: {e}", exc_info=True)
         gr.Error(f"Error processing input image: {e}")
-        return default_return # Use default return with correct seed
+        return default_return # Use default return structure
 
+    # Calculate intermediate dimensions for the model
     w_proc, h_proc = processed_input_image.size
     control_image_w = w_proc * INTERNAL_PROCESSING_FACTOR
     control_image_h = h_proc * INTERNAL_PROCESSING_FACTOR
 
-    # Failsafe clamp (remains the same)
-    if control_image_w * control_image_h > MAX_PIXEL_BUDGET * 1.05:
+    # Failsafe clamp if budget is still somehow exceeded after input processing
+    if control_image_w * control_image_h > MAX_PIXEL_BUDGET * 1.05: # Add 5% margin just in case
         scale_factor = (MAX_PIXEL_BUDGET / (control_image_w * control_image_h)) ** 0.5
-        control_image_w = max(8, int(control_image_w * scale_factor)); control_image_w -= control_image_w % 8
-        control_image_h = max(8, int(control_image_h * scale_factor)); control_image_h -= control_image_h % 8
-        logging.warning(f"Control image dims clamped post-processing: {control_image_w}x{control_image_h}.")
+        control_image_w = max(8, int(control_image_w * scale_factor))
+        control_image_h = max(8, int(control_image_h * scale_factor))
+        # Ensure multiple of 8 after scaling
+        control_image_w = max(8, control_image_w - control_image_w % 8)
+        control_image_h = max(8, control_image_h - control_image_h % 8)
+        logging.warning(f"Control image dimensions clamped post-processing to {control_image_w}x{control_image_h} to fit budget.")
         gr.Warning(f"Control image dimensions further clamped to {control_image_w}x{control_image_h}.")
 
+    # Prepare control image (resized input for ControlNet)
     logging.info(f"Resizing processed input {w_proc}x{h_proc} to control image {control_image_w}x{control_image_h}")
     try:
         control_image = processed_input_image.resize((control_image_w, control_image_h), Image.Resampling.LANCZOS)
     except ValueError as resize_err:
-         logging.error(f"Error resizing to control image: {resize_err}")
+         logging.error(f"Error resizing processed input to control image: {resize_err}")
          gr.Error(f"Failed to prepare control image: {resize_err}")
-         return default_return # Use default return
+         return default_return
 
+    # Setup generator for reproducibility
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    gr.Info(f"Generating intermediate image ({control_image_w}x{control_image_h}, {num_inference_steps} steps)...")
+    # --- Run the Diffusion Pipeline ---
+    gr.Info(f"Generating intermediate image at {INTERNAL_PROCESSING_FACTOR}x quality ({control_image_w}x{control_image_h}) with {num_inference_steps} steps...")
     logging.info(f"Running pipeline: size={control_image_w}x{control_image_h}, steps={num_inference_steps}")
-    intermediate_result_image = None # Initialize
+    intermediate_result_image = None # Initialize to None
     try:
         with torch.inference_mode():
-             # Progress bar integration can be added here if needed
+            # The actual model inference call
             intermediate_result_image = pipe(
-                prompt="", control_image=control_image,
+                prompt="", # No text prompt needed for this upscaler
+                control_image=control_image,
                 controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-                num_inference_steps=num_inference_steps, guidance_scale=0.0,
-                height=control_image_h, width=control_image_w, generator=generator,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0, # Guidance scale typically 0 for ControlNet-only tasks
+                height=control_image_h, # Target height for the model
+                width=control_image_w,   # Target width for the model
+                generator=generator,
+                # Can add callback for step progress if needed:
+                # callback_on_step_end=lambda step, t, latents: progress(step / num_inference_steps)
             ).images[0]
-        logging.info(f"Pipeline finished. Intermediate size: {intermediate_result_image.size if intermediate_result_image else 'None'}")
+        logging.info(f"Pipeline execution finished. Intermediate image size: {intermediate_result_image.size if intermediate_result_image else 'None'}")
 
-    # --- Catch specific errors if needed, otherwise general Exception ---
     except torch.cuda.OutOfMemoryError as oom_error:
-         logging.error(f"OOM during pipeline: {oom_error}", exc_info=True)
-         gr.Error(f"OOM generating {control_image_w}x{control_image_h}. Try smaller input/factor.")
-         if device == 'cuda': torch.cuda.empty_cache()
-         return default_return # Use default return
-    except Exception as e: # Catches the torch.compile error too
-        logging.error(f"Error during pipeline execution: {e}", exc_info=True) # Log full traceback
-        # Provide a more specific error message if it's the known compile issue
+         # Handle specific OOM error
+         logging.error(f"CUDA Out of Memory during pipeline execution: {oom_error}", exc_info=True)
+         gr.Error(f"Ran out of GPU memory trying to generate intermediate {control_image_w}x{control_image_h}. Try reducing the Final Upscale Factor or using a smaller input image.")
+         if device == 'cuda': torch.cuda.empty_cache() # Try to clear cache
+         return default_return
+    except Exception as e:
+        # Handle other pipeline errors, including potential torch.compile issues
+        logging.error(f"Error during pipeline execution: {e}", exc_info=True)
+        # Check if it looks like the known compile error
         if "dynamic shape operator" in str(e) or "Unsupported" in str(e.__class__):
-             gr.Error(f"Inference failed: torch.compile issue encountered ({type(e).__name__}). Try restarting the Space or disabling compilation if persistent.")
+             gr.Error(f"Inference failed: torch.compile issue encountered ({type(e).__name__}). The model attempted to run unoptimized. If this persists, compilation might need to be disabled in the code.")
         else:
-            gr.Error(f"Inference failed: {e}")
-        return default_return # Use default return
+             gr.Error(f"Inference failed: {e}")
+        return default_return
 
-    # --- Check if intermediate image was actually created ---
+    # --- Post-Pipeline Checks and Resizing ---
     if not intermediate_result_image:
-         logging.error("Intermediate result is None after pipeline (but no exception caught).")
+         # Should ideally not happen if no exception was caught, but check anyway
+         logging.error("Intermediate result image is None after pipeline execution without exception.")
          gr.Error("Inference produced no result image unexpectedly.")
-         return default_return # Use default return
+         return default_return
 
-    # --- Final Resizing (remains the same logic) ---
+    # Calculate final target dimensions based on ORIGINAL input size and FINAL upscale factor
     if was_input_resized:
+        # Base final size on the downscaled input that was processed
         final_target_w = w_proc * final_upscale_factor
         final_target_h = h_proc * final_upscale_factor
-        logging.warning(f"Input downscaled. Final size based on processed: {w_proc}x{h_proc}*{final_upscale_factor}x -> {final_target_w}x{final_target_h}")
-        gr.Info(f"Input downscaled. Final size approx {final_target_w}x{final_target_h}.")
+        logging.warning(f"Input was downscaled. Final size based on processed input: {w_proc}x{h_proc} * {final_upscale_factor}x -> {final_target_w}x{final_target_h}")
+        gr.Info(f"Input was downscaled. Final size target approx {final_target_w}x{final_target_h}.")
     else:
+        # Base final size on the original input size
         final_target_w = w_original * final_upscale_factor
         final_target_h = h_original * final_upscale_factor
 
+    # Perform final resize from intermediate to target size
     final_result_image = intermediate_result_image
     current_w, current_h = intermediate_result_image.size
 
     if (current_w, current_h) != (final_target_w, final_target_h):
-        logging.info(f"Resizing intermediate {current_w}x{current_h} to final {final_target_w}x{final_target_h}")
+        logging.info(f"Resizing intermediate image from {current_w}x{current_h} to final target size {final_target_w}x{final_target_h} (using {final_upscale_factor}x factor)")
         gr.Info(f"Resizing from intermediate {current_w}x{current_h} to final {final_target_w}x{final_target_h}...")
         try:
+            # Ensure target dimensions are valid before resizing
             if final_target_w > 0 and final_target_h > 0:
+                 # Use high-quality LANCZOS for downsampling or general resizing
                  final_result_image = intermediate_result_image.resize((final_target_w, final_target_h), Image.Resampling.LANCZOS)
             else:
-                 gr.Warning(f"Invalid final target ({final_target_w}x{final_target_h}). Skipping resize.")
-                 final_result_image = intermediate_result_image
+                 # Avoid resizing if target dimensions are invalid
+                 gr.Warning(f"Invalid final target dimensions ({final_target_w}x{final_target_h}). Skipping final resize.")
+                 final_result_image = intermediate_result_image # Keep intermediate
         except Exception as resize_e:
-            logging.error(f"Could not resize final image: {resize_e}")
-            gr.Warning(f"Failed final resize. Returning intermediate {current_w}x{current_h}.")
-            final_result_image = intermediate_result_image
+            # Handle potential errors during final resize
+            logging.error(f"Could not resize intermediate image to final size: {resize_e}")
+            gr.Warning(f"Failed to resize to final {final_upscale_factor}x. Returning intermediate {INTERNAL_PROCESSING_FACTOR}x result ({current_w}x{current_h}).")
+            final_result_image = intermediate_result_image # Fallback to intermediate
     else:
-        logging.info("Intermediate size matches final target. No final resize needed.")
+        # No resize needed if intermediate matches final target
+        logging.info(f"Intermediate size {current_w}x{current_h} matches final target size. No final resize needed.")
 
     logging.info(f"Inference successful. Final output size: {final_result_image.size}")
 
-    # --- Base64 Encoding (remains the same) ---
+    # --- Base64 Encoding for API output ---
     base64_string = None
     if final_result_image:
         try:
-            buffered = io.BytesIO(); final_result_image.save(buffered, format="WEBP", quality=90)
+            buffered = io.BytesIO()
+            # Save as WEBP for potentially smaller size, adjust quality as needed
+            final_result_image.save(buffered, format="WEBP", quality=90)
             img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            # Format as a data URL
             base64_string = f"data:image/webp;base64,{img_str}"
-            logging.info(f"Encoded result to Base64 (len: {len(base64_string)}).")
+            logging.info(f"Encoded result image to Base64 string (length: {len(base64_string)} chars).")
         except Exception as enc_err:
-             logging.error(f"Failed Base64 encoding: {enc_err}", exc_info=True)
+             # Log if encoding fails but don't stop the process
+             logging.error(f"Failed to encode result image to Base64: {enc_err}", exc_info=True)
+             base64_string = None # Ensure it's None if encoding failed
 
-    # --- SUCCESS RETURN ---
+    # --- Return results for Gradio ---
     return [[original_input_pil, final_result_image], seed, base64_string]
 
-# --- Gradio Interface (Gradio UI Definition) ---
-# ... (Gradio definition remains the same, ensure inputs/outputs match infer) ...
+
+# --- Gradio Interface Definition ---
 with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as demo:
-    gr.Markdown(f"""
+    gr.Markdown(
+        f"""
     # ⚡ Flux.1-dev Upscaler ControlNet ⚡
-    Upscale images using Flux.1-dev Upscaler ControlNet on **{power_device}**.
-    Internal processing at **{INTERNAL_PROCESSING_FACTOR}x** quality, resized to **Final Upscale Factor**.
-    **Speed Up:** Reduce `Inference Steps` (try 10-15). `torch.compile` enabled (may fail, see logs).
-    *Limit*: ~**{MAX_PIXEL_BUDGET/1_000_000:.1f} megapixels** intermediate size.
-    """)
+    Upscale images using the [Flux.1-dev Upscaler ControlNet](https://huggingface.co/jasperai/Flux.1-dev-Controlnet-Upscaler) model based on [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev).
+    Currently running on **{power_device}**. Hardware provided by Hugging Face 🤗.
+
+    **How it works:** This demo uses an internal processing scale of **{INTERNAL_PROCESSING_FACTOR}x** for potentially higher detail generation (slower),
+    then resizes the result to your selected **Final Upscale Factor**.
+
+    **To Speed Up:**
+    1.  **Reduce `Inference Steps`:** Fewer steps = faster generation (try 10-15).
+    2.  **(Code Change Needed):** Reduce `INTERNAL_PROCESSING_FACTOR` in the script (e.g., to 3) for direct computation reduction (may lower detail).
+    3.  `torch.compile` is attempted (`mode="default"`) which might provide speedup after the first run (check logs for success/failure).
+
+    *Limit*: Intermediate processing resolution capped around **{MAX_PIXEL_BUDGET/1_000_000:.1f} megapixels** ({int(MAX_PIXEL_BUDGET**0.5)}x{int(MAX_PIXEL_BUDGET**0.5)}).
+    """
+    )
+
     with gr.Row():
-        with gr.Column(scale=2): input_im = gr.Image(label="Input Image", type="pil", height=350, sources=["upload", "clipboard"])
+        with gr.Column(scale=2):
+            input_im = gr.Image(
+                label="Input Image",
+                type="pil",
+                height=350,
+                sources=["upload", "clipboard"],
+            )
         with gr.Column(scale=1):
-            upscale_factor_slider = gr.Slider(label="Final Upscale Factor", info=f"Output size. Internal uses {INTERNAL_PROCESSING_FACTOR}x quality.", minimum=1, maximum=INTERNAL_PROCESSING_FACTOR, step=1, value=min(2, INTERNAL_PROCESSING_FACTOR))
-            num_inference_steps = gr.Slider(label="Inference Steps", info="Fewer=faster (try 10-15).", minimum=4, maximum=50, step=1, value=15)
-            controlnet_conditioning_scale = gr.Slider(label="ControlNet Scale", info="Guidance strength.", minimum=0.0, maximum=1.5, step=0.05, value=0.6)
+            upscale_factor_slider = gr.Slider(
+                label="Final Upscale Factor",
+                info=f"Output size relative to input. Internal processing uses {INTERNAL_PROCESSING_FACTOR}x quality.",
+                minimum=1,
+                maximum=INTERNAL_PROCESSING_FACTOR, # Max limited by internal factor
+                step=1,
+                value=min(2, INTERNAL_PROCESSING_FACTOR) # Default to 2x or internal factor if smaller
+            )
+            num_inference_steps = gr.Slider(
+                label="Inference Steps",
+                info="Fewer steps = faster. Try 10-15.",
+                minimum=4, maximum=50, step=1, value=15 # Defaulting to 15 for speed
+            )
+            controlnet_conditioning_scale = gr.Slider(
+                label="ControlNet Conditioning Scale",
+                info="Strength of ControlNet guidance.",
+                minimum=0.0, maximum=1.5, step=0.05, value=0.6
+            )
             with gr.Row():
                  seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                  randomize_seed = gr.Checkbox(label="Random", value=True, scale=0, min_width=80)
             run_button = gr.Button("⚡ Upscale Image", variant="primary", scale=1)
-    with gr.Row(): result_slider = ImageSlider(label="Input / Output Comparison", type="pil", interactive=False, show_label=True, position=0.5)
+
+    with gr.Row():
+        result_slider = ImageSlider(
+            label="Input / Output Comparison",
+            type="pil",
+            interactive=False, # Output only
+            show_label=True,
+            position=0.5 # Start slider in the middle
+            )
+
     output_seed = gr.Textbox(label="Seed Used", interactive=False, visible=True, scale=1)
+    # Hidden output for API usage
     api_base64_output = gr.Textbox(label="API Base64 Output", interactive=False, visible=False)
 
-    example_dir = "examples"; example_files = ["image_2.jpg", "image_4.jpg", "low_res_face.png", "low_res_landscape.png"]
+    # --- Examples ---
+    example_dir = "examples"
+    example_files = ["image_2.jpg", "image_4.jpg", "low_res_face.png", "low_res_landscape.png"]
     example_paths = [os.path.join(example_dir, f) for f in example_files if os.path.exists(os.path.join(example_dir, f))]
+
     if example_paths:
         gr.Examples(
+            # Examples use the new defaults: final factor 2x (or less), steps 15
             examples=[ [path, min(2, INTERNAL_PROCESSING_FACTOR), 15, 0.6, random.randint(0,MAX_SEED), True] for path in example_paths ],
             inputs=[ input_im, upscale_factor_slider, num_inference_steps, controlnet_conditioning_scale, seed, randomize_seed, ],
-            outputs=[result_slider, output_seed], fn=infer, cache_examples="lazy",
-            label=f"Examples (Click: {min(2, INTERNAL_PROCESSING_FACTOR)}x Output, 15 Steps)", run_on_click=True)
-    else: gr.Markdown(f"*No example images found in '{example_dir}'.*")
-    gr.Markdown("---"); gr.Markdown("**Disclaimer:** For illustrative purposes.")
-
-    run_button.click(fn=infer,
-        inputs=[seed, randomize_seed, input_im, num_inference_steps, upscale_factor_slider, controlnet_conditioning_scale],
-        outputs=[result_slider, output_seed, api_base64_output], api_name="upscale")
-
+            # Map outputs to UI components; base64 ignored by Examples UI
+            outputs=[result_slider, output_seed],
+            fn=infer, # Function to call when example is clicked
+            cache_examples="lazy", # Cache results for examples
+            label=f"Example Images (Click to Run with {min(2, INTERNAL_PROCESSING_FACTOR)}x Output, 15 Steps)",
+            run_on_click=True
+        )
+    else:
+        gr.Markdown(f"*No example images found in '{example_dir}' directory.*")
+
+    gr.Markdown("---")
+    gr.Markdown("**Disclaimer:** Demo for illustrative purposes. Users are responsible for generated content.")
+
+    # Connect button click to the inference function
+    run_button.click(
+        fn=infer,
+        inputs=[
+            seed,
+            randomize_seed,
+            input_im,
+            num_inference_steps,
+            upscale_factor_slider,
+            controlnet_conditioning_scale,
+        ],
+        # Map all return values from infer to the correct output components
+        outputs=[result_slider, output_seed, api_base64_output],
+        api_name="upscale" # Define API endpoint name
+    )
+
+# Launch the Gradio app
+# queue manages concurrent users/requests
+# share=False means no public link generated by default
 demo.queue(max_size=10).launch(share=False, show_api=True)