Update app.py

app.py

@@ -6,14 +6,12 @@ import io
 import base64
 import gradio as gr
 import numpy as np
-import spaces # Ensure spaces is imported for the decorator
+import spaces
 import torch
-# --- Optional: Add this if trying Solution 2 (suppress_errors) later ---
-# import torch._dynamo
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
-from gradio_imageslider import ImageSlider
-from PIL import Image, ImageOps
+from gradio_imageslider import ImageSlider # Ensure this is installed: pip install gradio_imageslider
+from PIL import Image, ImageOps # Import ImageOps for exif transpose
 from huggingface_hub import snapshot_download
 
 # --- Setup Logging and Device ---
@@ -21,25 +19,34 @@ logging.basicConfig(level=logging.INFO)
 warnings.filterwarnings("ignore")
 
 css = """
-#col-container { margin: 0 auto; max-width: 512px; }
-.gradio-container { max-width: 900px !important; margin: auto !important; }
+#col-container {
+    margin: 0 auto;
+    max-width: 512px; /* Increased max-width slightly for better layout */
+}
+.gradio-container {
+    max-width: 900px !important; /* Control overall container width */
+    margin: auto !important;
+}
 """
 
 if torch.cuda.is_available():
     power_device = "GPU"
     device = "cuda"
-    torch_dtype = torch.bfloat16
+    torch_dtype = torch.bfloat16 # Use bfloat16 for GPU for better performance/memory
 else:
     power_device = "CPU"
     device = "cpu"
-    torch_dtype = torch.float32
+    torch_dtype = torch.float32 # Use float32 for CPU
+
 logging.info(f"Selected device: {device} | Data type: {torch_dtype}")
 
 # --- Authentication and Model Download ---
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+
+# Define model IDs
 flux_model_id = "black-forest-labs/FLUX.1-dev"
 controlnet_model_id = "jasperai/Flux.1-dev-Controlnet-Upscaler"
-local_model_dir = flux_model_id.split('/')[-1]
+local_model_dir = flux_model_id.split('/')[-1] # e.g., "FLUX.1-dev"
 pipe = None
 
 try:
@@ -49,7 +56,7 @@ try:
         repo_type="model",
         ignore_patterns=["*.md", "*.gitattributes"],
         local_dir=local_model_dir,
-        token=huggingface_token
+        token=huggingface_token,
     )
     logging.info(f"Base model downloaded/verified in: {model_path}")
 
@@ -68,88 +75,60 @@ try:
     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"):
-        # --- Using "default" mode as first attempt to fix compile errors ---
-        compile_mode = "default"
-        # --- 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
-        # --- End Alternative ---
-
-        logging.info(f"Attempting to compile the pipeline transformer with torch.compile (mode='{compile_mode}')...")
-        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 for pre-compilation
-        except Exception as e:
-            logging.warning(f"torch.compile failed (mode='{compile_mode}'): {e}. Running unoptimized.")
-            # --- 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 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}")
+    logging.error(f"FATAL: Error during model loading: {e}", exc_info=True)
+    # --- Simplified Error Handling for Brevity ---
+    print(f"FATAL ERROR DURING MODEL LOAD: {e}")
+    raise SystemExit(f"Model loading failed: {e}")
+
 
 # --- Constants ---
 MAX_SEED = 2**32 - 1
-MAX_PIXEL_BUDGET = 1280 * 1280 # Max pixels for the *intermediate* image
-INTERNAL_PROCESSING_FACTOR = 4 # Factor used for quality generation
+MAX_PIXEL_BUDGET = 1280 * 1280
+# --- NEW: Define the internal factor for quality ---
+INTERNAL_PROCESSING_FACTOR = 4
 
-# --- Image Processing Function (process_input) ---
+# --- Image Processing Function (Modified) ---
 def process_input(input_image):
-    """Processes input image: handles orientation, converts to RGB, checks budget, rounds dimensions."""
+    """Processes the input image for the pipeline.
+       The pixel budget check uses the fixed INTERNAL_PROCESSING_FACTOR."""
     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)
-        # 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:
-        # 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 width or height.")
 
-    # Calculate target intermediate size based on INTERNAL factor for budget check
+    # Calculate target 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() # Work on a copy
+    input_image_to_process = input_image.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)) # Ensure min size 8
-            input_h_resized = max(8, int(h * input_scale_factor)) # Ensure min size 8
+            input_w_resized = int(w * input_scale_factor)
+            input_h_resized = int(h * input_scale_factor)
+            # Ensure minimum size of 8x8
+            input_w_resized = max(8, input_w_resized)
+            input_h_resized = max(8, input_h_resized)
             intermediate_w = input_w_resized * INTERNAL_PROCESSING_FACTOR
             intermediate_h = input_h_resized * INTERNAL_PROCESSING_FACTOR
 
@@ -162,12 +141,12 @@ def process_input(input_image):
                 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 # Flag that original dimensions were lost for final scaling
+            was_resized = True # Flag that original dimensions are lost for precise final scaling
 
-    # Round processed input dimensions down to nearest multiple of 8 (required by some models)
+    # Round processed input dimensions to be multiple of 8
     w_proc, h_proc = input_image_to_process.size
-    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
+    w_final_proc = max(8, w_proc - w_proc % 8)
+    h_final_proc = max(8, h_proc - h_proc % 8)
 
     if (w_proc, h_proc) != (w_final_proc, h_final_proc):
         logging.info(f"Rounding processed input dimensions from {w_proc}x{h_proc} to {w_final_proc}x{h_final_proc}")
@@ -175,147 +154,120 @@ def process_input(input_image):
 
     return input_image_to_process, w_original, h_original, was_resized
 
-
-# --- Inference Function (infer) ---
-# --- MODIFIED GPU DURATION ---
+# --- Inference Function (Modified) ---
 @spaces.GPU(duration=75)
 def infer(
     seed,
     randomize_seed,
     input_image,
     num_inference_steps,
-    final_upscale_factor,
+    final_upscale_factor, # Renamed for clarity internally
     controlnet_conditioning_scale,
-    progress=gr.Progress(track_tqdm=True), # Gradio progress tracking
+    progress=gr.Progress(track_tqdm=True),
 ):
-    """Runs the Flux ControlNet upscaling pipeline."""
     global pipe
-    # --- 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
+        return [[None, None], 0, None]
 
-    original_input_pil = input_image # Keep reference to original input
+    original_input_pil = input_image # Keep ref for slider
 
-    # Handle missing input image
     if input_image is None:
          gr.Warning("Please provide an input image.")
-         # 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
+         return [[None, None], seed or 0, None]
 
-    # 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
-    default_return[0][0] = original_input_pil
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
 
-    # Ensure numerical inputs are integers
+    # Ensure final_upscale_factor is an integer
     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
+        gr.Warning(f"Selected upscale factor ({final_upscale_factor}x) is larger than internal processing factor ({INTERNAL_PROCESSING_FACTOR}x). "
+                   f"Results might not be optimal. Clamping final factor to {INTERNAL_PROCESSING_FACTOR}x for this run.")
+        final_upscale_factor = INTERNAL_PROCESSING_FACTOR # Prevent upscaling *beyond* internal processing
 
     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}"
+        f"Starting inference with seed: {seed}, "
+        f"Internal Processing Factor: {INTERNAL_PROCESSING_FACTOR}x, "
+        f"Final Output Factor: {final_upscale_factor}x, "
+        f"Steps: {num_inference_steps}, CNet Scale: {controlnet_conditioning_scale}"
     )
 
-    # Process the input image
     try:
+        # process_input now implicitly uses INTERNAL_PROCESSING_FACTOR for budget checks
         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 structure
+        return [[original_input_pil, None], seed, None]
 
-    # Calculate intermediate dimensions for the model
     w_proc, h_proc = processed_input_image.size
+
+    # Calculate control image dimensions using INTERNAL_PROCESSING_FACTOR
     control_image_w = w_proc * INTERNAL_PROCESSING_FACTOR
     control_image_h = h_proc * INTERNAL_PROCESSING_FACTOR
 
-    # 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
+    # Clamp control image size if it *still* exceeds budget (e.g., due to rounding or small inputs)
+    # This check should technically be redundant if process_input worked correctly, but good failsafe.
+    if control_image_w * control_image_h > MAX_PIXEL_BUDGET * 1.05: # Add a small margin
         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_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.")
+        logging.warning(f"Control image dimensions clamped to {control_image_w}x{control_image_h} post-processing 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}")
+    logging.info(f"Resizing processed input {w_proc}x{h_proc} to control image {control_image_w}x{control_image_h} (using {INTERNAL_PROCESSING_FACTOR}x factor)")
     try:
+        # Use the processed input image for control, resized to the intermediate size
         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 processed input to control image: {resize_err}")
          gr.Error(f"Failed to prepare control image: {resize_err}")
-         return default_return
+         return [[original_input_pil, None], seed, None]
 
-    # Setup generator for reproducibility
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # --- 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 to None
+    # --- Run the Pipeline at INTERNAL_PROCESSING_FACTOR scale ---
+    gr.Info(f"Generating intermediate image at {INTERNAL_PROCESSING_FACTOR}x quality ({control_image_w}x{control_image_h})...")
+    logging.info(f"Running pipeline with size: {control_image_w}x{control_image_h}")
+    intermediate_result_image = None
     try:
         with torch.inference_mode():
-            # The actual model inference call
             intermediate_result_image = pipe(
-                prompt="", # No text prompt needed for this upscaler
-                control_image=control_image,
+                prompt="",
+                control_image=control_image, # Control image IS the intermediate size
                 controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-                num_inference_steps=num_inference_steps,
-                guidance_scale=0.0, # Guidance scale typically 0 for ControlNet-only tasks
+                num_inference_steps=int(num_inference_steps),
+                guidance_scale=0.0,
                 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 execution finished. Intermediate image size: {intermediate_result_image.size if intermediate_result_image else 'None'}")
 
     except torch.cuda.OutOfMemoryError as oom_error:
-         # 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
+         gr.Error(f"Ran out of GPU memory trying to generate intermediate {control_image_w}x{control_image_h}.")
+         if device == 'cuda': torch.cuda.empty_cache()
+         return [[original_input_pil, None], seed, None]
     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__}). 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
-
-    # --- Post-Pipeline Checks and Resizing ---
+        gr.Error(f"Inference failed: {e}")
+        return [[original_input_pil, None], seed, None]
+
     if not intermediate_result_image:
-         # 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
+         logging.error("Intermediate result image is None after pipeline execution.")
+         gr.Error("Inference produced no result image.")
+         return [[original_input_pil, None], seed, None]
 
+    # --- Final Resizing to User's Desired Scale ---
     # Calculate final target dimensions based on ORIGINAL input size and FINAL upscale factor
+    # If input was resized, we scale the *processed* input size instead, as original is unknown
     if was_input_resized:
         # Base final size on the downscaled input that was processed
         final_target_w = w_proc * final_upscale_factor
@@ -327,54 +279,48 @@ def infer(
         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
 
+    # Only resize if the intermediate size doesn't match the final desired size
     if (current_w, current_h) != (final_target_w, 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
+                 # Use LANCZOS for downsampling, it's high quality
                  final_result_image = intermediate_result_image.resize((final_target_w, final_target_h), Image.Resampling.LANCZOS)
             else:
-                 # 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:
-            # 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:
-        # 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 for API output ---
+    # --- Base64 Encoding (No change needed here) ---
     base64_string = None
     if final_result_image:
         try:
             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 image to Base64 string (length: {len(base64_string)} chars).")
         except Exception as enc_err:
-             # 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
 
-    # --- Return results for Gradio ---
+    # Return original input and the FINAL processed image
     return [[original_input_pil, final_result_image], seed, base64_string]
 
 
-# --- Gradio Interface Definition ---
+# --- Gradio Interface (Minor Text Updates) ---
 with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as demo:
     gr.Markdown(
         f"""
@@ -382,15 +328,11 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
     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**.
+    **How it works:** This demo uses an internal processing scale of **{INTERNAL_PROCESSING_FACTOR}x** for potentially higher detail generation,
+    then resizes the result to your selected **Final Upscale Factor**. This aims for {INTERNAL_PROCESSING_FACTOR}x quality at your desired output resolution.
 
-    **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)}).
+    *Note*: Intermediate processing resolution is limited to approximately **{MAX_PIXEL_BUDGET/1_000_000:.1f} megapixels** ({int(MAX_PIXEL_BUDGET**0.5)}x{int(MAX_PIXEL_BUDGET**0.5)}) due to resource constraints.
+    The *diffusion process time* is mainly determined by this intermediate size, not the final output size.
     """
     )
 
@@ -403,24 +345,10 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
                 sources=["upload", "clipboard"],
             )
         with gr.Column(scale=1):
-            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
-            )
+            # Renamed slider label for clarity
+            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, step=1, value=2) # Default to 2x, max is now INTERNAL_PROCESSING_FACTOR
+            num_inference_steps = gr.Slider(label="Inference Steps", minimum=4, maximum=50, step=1, value=15)
+            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)
@@ -430,30 +358,29 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
         result_slider = ImageSlider(
             label="Input / Output Comparison",
             type="pil",
-            interactive=False, # Output only
+            interactive=False,
             show_label=True,
-            position=0.5 # Start slider in the middle
+            position=0.5
             )
 
     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)
 
-    # --- Examples ---
+    # --- Examples (Updated default factor if needed) ---
     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 ],
+            # Examples now use the new default of 2x for the final factor
+            examples=[ [path, 2, 15, 0.6, random.randint(0,MAX_SEED), True] for path in example_paths ],
+            # Ensure inputs match the order expected by `infer` now
             inputs=[ input_im, upscale_factor_slider, num_inference_steps, controlnet_conditioning_scale, seed, randomize_seed, ],
-            # 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)",
+            outputs=[result_slider, output_seed], # Base64 output ignored by Examples
+            fn=infer,
+            cache_examples="lazy",
+            label="Example Images (Click to Run with 2x Output)",
             run_on_click=True
         )
     else:
@@ -462,7 +389,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
     gr.Markdown("---")
     gr.Markdown("**Disclaimer:** Demo for illustrative purposes. Users are responsible for generated content.")
 
-    # Connect button click to the inference function
+    # Connect button click
     run_button.click(
         fn=infer,
         inputs=[
@@ -470,15 +397,12 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
             randomize_seed,
             input_im,
             num_inference_steps,
-            upscale_factor_slider,
+            upscale_factor_slider, # Use the slider value here
             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
+        api_name="upscale"
     )
 
 # 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)