Update app.py

app.py

@@ -1,324 +1,253 @@
+# ---- Imports ----
 import logging
 import random
 import warnings
 import os
-import io # Ensure io is imported
-import base64 # Ensure base64 is imported
+import io
+import base64
 import gradio as gr
 import numpy as np
 import spaces
 import torch
+# --- Add this if you want to try Solution 2 later ---
+# import torch._dynamo
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
-from gradio_imageslider import ImageSlider # Ensure this is installed: pip install gradio_imageslider
-from PIL import Image, ImageOps # Import ImageOps for exif transpose
+from gradio_imageslider import ImageSlider
+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")
 
 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; }
+.gradio-container { max-width: 900px !important; margin: auto !important; }
 """
 
 if torch.cuda.is_available():
-    power_device = "GPU"
-    device = "cuda"
-    torch_dtype = torch.bfloat16 # Use bfloat16 for GPU for better performance/memory
+    power_device = "GPU"; device = "cuda"; torch_dtype = torch.bfloat16
 else:
-    power_device = "CPU"
-    device = "cpu"
-    torch_dtype = torch.float32 # Use float32 for CPU
-
+    power_device = "CPU"; device = "cpu"; torch_dtype = torch.float32
 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] # e.g., "FLUX.1-dev"
+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"):
-        logging.info("Attempting to compile the pipeline transformer with torch.compile...")
+        # --- TRY THIS FIRST: Change mode to 'default' ---
+        compile_mode = "default"
+        # --- Alternative (Solution 2): Uncomment these lines ---
+        # 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:
-            # Modes: 'default', 'reduce-overhead', 'max-autotune'
-            # 'max-autotune' might give best runtime performance but takes longer to compile initially
-            pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
-            # You could potentially compile other components too, but start with the transformer
-            # pipe.controlnet = torch.compile(pipe.controlnet, mode="max-autotune", fullgraph=True)
+            pipe.transformer = torch.compile(pipe.transformer, mode=compile_mode, fullgraph=True)
             logging.info("Pipeline transformer compiled successfully.")
-            # Optional: Add a dummy inference run here to trigger compilation during startup
-            # This avoids the compilation delay on the *first* user request.
-            # try:
-            #     logging.info("Running dummy inference to finalize compilation...")
-            #     _ = pipe(prompt="", control_image=Image.new('RGB', (64, 64)), height=64*4, width=64*4, num_inference_steps=1, guidance_scale=0.0, output_type="latent")
-            #     logging.info("Dummy inference completed.")
-            # except Exception as compile_run_e:
-            #     logging.warning(f"Dummy inference after compile failed (might be ok): {compile_run_e}")
+            # Optional dummy inference run can go here
         except Exception as e:
-            logging.warning(f"torch.compile failed: {e}. Running unoptimized.")
+            logging.warning(f"torch.compile failed (mode='{compile_mode}'): {e}. Running unoptimized.")
+            # --- 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.")
 
-    # --- OPTIMIZATION: Consider xformers ---
-    # If torch.compile doesn't provide enough speedup or isn't available,
-    # you can try installing and enabling xformers.
-    # 1. Add `xformers` to your requirements.txt or install it (`pip install xformers`).
-    # 2. Uncomment and add this code block *before* the torch.compile block:
-    # try:
-    #     import xformers
-    #     pipe.enable_xformers_memory_efficient_attention()
-    #     logging.info("Enabled xformers memory efficient attention.")
-    # except ImportError:
-    #     logging.info("xformers not installed. Skipping.")
-    # except Exception as e:
-    #     logging.warning(f"Could not enable xformers: {e}.")
+    # --- (Optional xformers code would go here if used) ---
 
     logging.info("Pipeline ready for inference.")
 
-
 except Exception as e:
     logging.error(f"FATAL: Error during model loading or setup: {e}", exc_info=True)
-    # Simple error display if Gradio Blocks object isn't ready
     print(f"FATAL ERROR DURING MODEL LOAD/SETUP: {e}")
-    # You might want to use the Gradio error block structure here if `gr` is available
-    # with gr.Blocks() as demo_error: ... etc.
     raise SystemExit(f"Model loading/setup failed: {e}")
 
 
 # --- Constants ---
 MAX_SEED = 2**32 - 1
-MAX_PIXEL_BUDGET = 1280 * 1280 # Max pixels for the *intermediate* high-res image
-
-# --- SPEED VS QUALITY ---
-# INTERNAL_PROCESSING_FACTOR: Determines the scale the diffusion model *targets*
-# Higher values (like 4) aim for more detail generation but are slower.
-# Lower values (like 3 or 2) will be faster but might produce less detail enhancement.
-# You were aiming for 4x quality, so we keep it at 4. Reducing this is a direct speedup trade-off.
+MAX_PIXEL_BUDGET = 1280 * 1280
 INTERNAL_PROCESSING_FACTOR = 4
 
-# --- Image Processing Function (Uses INTERNAL_PROCESSING_FACTOR for budgeting) ---
+# --- Image Processing Function (process_input) ---
+# ... (process_input function remains the same) ...
 def process_input(input_image):
-    """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!")
+    if input_image is None: raise gr.Error("Input image is missing!")
     try:
         input_image = ImageOps.exif_transpose(input_image)
-        if input_image.mode != 'RGB':
-            logging.info(f"Converting input image from {input_image.mode} to RGB")
-            input_image = input_image.convert('RGB')
+        if input_image.mode != 'RGB': input_image = input_image.convert('RGB')
         w, h = input_image.size
-    except AttributeError:
-        raise gr.Error("Invalid input image format. Please provide a valid image file.")
-    except Exception as img_err:
-         raise gr.Error(f"Could not process input image: {img_err}")
+    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}")
 
     w_original, h_original = w, h
-    if w == 0 or h == 0:
-        raise gr.Error("Input image has zero width or height.")
+    if w == 0 or h == 0: raise gr.Error("Input image has zero dimensions.")
 
-    # 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()
 
-    # Check if the *intermediate* size exceeds the budget
     if target_pixels_internal > MAX_PIXEL_BUDGET:
         max_input_pixels = MAX_PIXEL_BUDGET / (INTERNAL_PROCESSING_FACTOR**2)
         current_input_pixels = w * h
-
         if current_input_pixels > max_input_pixels:
             input_scale_factor = (max_input_pixels / current_input_pixels) ** 0.5
-            input_w_resized = int(w * input_scale_factor)
-            input_h_resized = int(h * input_scale_factor)
-            input_w_resized = max(8, input_w_resized) # Ensure min size
-            input_h_resized = max(8, input_h_resized) # Ensure min size
+            input_w_resized = max(8, int(w * input_scale_factor))
+            input_h_resized = max(8, int(h * input_scale_factor))
             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 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)}."
-            )
+            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)}.")
             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 are lost for precise final scaling
+            was_resized = True
 
-    # 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)
     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}")
+        logging.info(f"Rounding processed input dims {w_proc}x{h_proc} -> {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 (Runs pipeline at Internal Factor, resizes to Final Factor) ---
-@spaces.GPU(duration=70) # Keep GPU decorator
+
+# --- Inference Function (infer) ---
+@spaces.GPU(duration=180)
 def infer(
-    seed,
-    randomize_seed,
-    input_image,
-    num_inference_steps, # Reducing this is a direct way to speed up (quality trade-off)
-    final_upscale_factor, # User's desired final output scale
-    controlnet_conditioning_scale,
+    seed, randomize_seed, input_image, num_inference_steps,
+    final_upscale_factor, controlnet_conditioning_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
     global pipe
+    # --- IMPROVED ERROR HANDLING: Define default return early ---
+    default_return = [[input_image, None], int(seed) if seed is not None else 0, None]
+
     if pipe is None:
         gr.Error("Pipeline not loaded. Cannot perform inference.")
-        return [[None, None], 0, None]
+        return default_return # Use default return
 
-    original_input_pil = input_image
+    original_input_pil = input_image # Keep ref even if None initially
 
     if input_image is None:
          gr.Warning("Please provide an input image.")
-         return [[None, None], seed or 0, None]
+         # 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)
+    if randomize_seed: seed = random.randint(0, MAX_SEED)
     seed = int(seed)
+    # --- UPDATE DEFAULT RETURN SEED ---
+    default_return[1] = seed
+    # --- Ensure original image is in the default return ---
+    default_return[0][0] = original_input_pil
+
 
-    # Ensure factors are integers
     final_upscale_factor = int(final_upscale_factor)
-    num_inference_steps = int(num_inference_steps) # Ensure steps are int
+    num_inference_steps = int(num_inference_steps)
 
-    # Sanity check: final factor shouldn't exceed internal processing factor in this workflow
     if final_upscale_factor > INTERNAL_PROCESSING_FACTOR:
-        gr.Warning(f"Selected final upscale factor ({final_upscale_factor}x) is larger than internal processing factor ({INTERNAL_PROCESSING_FACTOR}x). "
-                   f"Clamping final factor to {INTERNAL_PROCESSING_FACTOR}x.")
+        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 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}"
-    )
+    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}")
 
     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 [[original_input_pil, None], seed, None]
+        return default_return # Use default return with correct seed
 
     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 exceeded (should be rare if process_input works)
-    if control_image_w * control_image_h > MAX_PIXEL_BUDGET * 1.05: # Small margin
+    # Failsafe clamp (remains the same)
+    if control_image_w * control_image_h > MAX_PIXEL_BUDGET * 1.05:
         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))
-        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.")
+        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}.")
         gr.Warning(f"Control image dimensions further clamped to {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)")
+    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 processed input to control image: {resize_err}")
+         logging.error(f"Error resizing to control image: {resize_err}")
          gr.Error(f"Failed to prepare control image: {resize_err}")
-         return [[original_input_pil, None], seed, None]
+         return default_return # Use default return
 
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # --- 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}) with {num_inference_steps} steps...")
-    logging.info(f"Running pipeline with size: {control_image_w}x{control_image_h}, steps: {num_inference_steps}")
-    intermediate_result_image = None
+    gr.Info(f"Generating intermediate image ({control_image_w}x{control_image_h}, {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
     try:
         with torch.inference_mode():
+             # Progress bar integration can be added here if needed
             intermediate_result_image = pipe(
-                prompt="",
-                control_image=control_image,
+                prompt="", control_image=control_image,
                 controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-                num_inference_steps=num_inference_steps, # Use the integer value
-                guidance_scale=0.0,
-                height=control_image_h,
-                width=control_image_w,
-                generator=generator,
-                # Add progress callback if desired, requires tqdm
-                # callback_on_step_end = ...
+                num_inference_steps=num_inference_steps, guidance_scale=0.0,
+                height=control_image_h, width=control_image_w, generator=generator,
             ).images[0]
-        logging.info(f"Pipeline execution finished. Intermediate image size: {intermediate_result_image.size if intermediate_result_image else 'None'}")
+        logging.info(f"Pipeline finished. Intermediate 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"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.")
+         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 [[original_input_pil, None], seed, None]
-    except Exception as e:
-        logging.error(f"Error during pipeline execution: {e}", exc_info=True)
-        gr.Error(f"Inference failed: {e}")
-        return [[original_input_pil, None], seed, None]
-
+         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
+        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.")
+        else:
+            gr.Error(f"Inference failed: {e}")
+        return default_return # Use default return
+
+    # --- Check if intermediate image was actually created ---
     if not intermediate_result_image:
-         logging.error("Intermediate result image is None after pipeline execution.")
-         gr.Error("Inference produced no result image.")
-         return [[original_input_pil, None], seed, None]
+         logging.error("Intermediate result is None after pipeline (but no exception caught).")
+         gr.Error("Inference produced no result image unexpectedly.")
+         return default_return # Use default return
 
-    # --- Final Resizing to User's Desired Scale ---
+    # --- Final Resizing (remains the same logic) ---
     if was_input_resized:
         final_target_w = w_proc * final_upscale_factor
         final_target_h = h_proc * final_upscale_factor
-        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}.")
+        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}.")
     else:
         final_target_w = w_original * final_upscale_factor
         final_target_h = h_original * final_upscale_factor
@@ -327,143 +256,74 @@ def infer(
     current_w, current_h = intermediate_result_image.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)")
+        logging.info(f"Resizing intermediate {current_w}x{current_h} to final {final_target_w}x{final_target_h}")
         gr.Info(f"Resizing from intermediate {current_w}x{current_h} to final {final_target_w}x{final_target_h}...")
         try:
             if final_target_w > 0 and final_target_h > 0:
                  final_result_image = intermediate_result_image.resize((final_target_w, final_target_h), Image.Resampling.LANCZOS)
             else:
-                 gr.Warning(f"Invalid final target dimensions ({final_target_w}x{final_target_h}). Skipping final resize.")
+                 gr.Warning(f"Invalid final target ({final_target_w}x{final_target_h}). Skipping resize.")
                  final_result_image = intermediate_result_image
         except Exception as resize_e:
-            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}).")
+            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
     else:
-        logging.info(f"Intermediate size {current_w}x{current_h} matches final target size. No final resize needed.")
+        logging.info("Intermediate size matches final target. No final resize needed.")
 
     logging.info(f"Inference successful. Final output size: {final_result_image.size}")
 
-    # --- Base64 Encoding ---
+    # --- Base64 Encoding (remains the same) ---
     base64_string = None
     if final_result_image:
         try:
-            buffered = io.BytesIO()
-            final_result_image.save(buffered, format="WEBP", quality=90) # WEBP is usually smaller
+            buffered = io.BytesIO(); final_result_image.save(buffered, format="WEBP", quality=90)
             img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
             base64_string = f"data:image/webp;base64,{img_str}"
-            logging.info(f"Encoded result image to Base64 string (length: {len(base64_string)} chars).")
+            logging.info(f"Encoded result to Base64 (len: {len(base64_string)}).")
         except Exception as enc_err:
-             logging.error(f"Failed to encode result image to Base64: {enc_err}", exc_info=True)
+             logging.error(f"Failed Base64 encoding: {enc_err}", exc_info=True)
 
+    # --- SUCCESS RETURN ---
     return [[original_input_pil, final_result_image], seed, base64_string]
 
-
-# --- Gradio Interface ---
+# --- Gradio Interface (Gradio UI Definition) ---
+# ... (Gradio definition remains the same, ensure inputs/outputs match infer) ...
 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 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**. This aims for {INTERNAL_PROCESSING_FACTOR}x quality at your desired output resolution.
-
-    **To Speed Up:**
-    1.  **Reduce `Inference Steps`:** Fewer steps = faster generation (potential quality decrease). Try 10-15 instead of 25.
-    2.  **(Code Change Needed):** Reduce `INTERNAL_PROCESSING_FACTOR` in the script (e.g., to 3). This directly reduces computation but may lower detail enhancement.
-    3.  `torch.compile` has been enabled (if using PyTorch 2.0+ on GPU) which should provide some speedup after the first run.
-
-    *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 and the number of steps.
-    """
-    )
-
+    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.
+    """)
     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):
-            # 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, # Max limited by internal factor
-                step=1,
-                value=min(2, INTERNAL_PROCESSING_FACTOR) # Default to 2x or internal factor if smaller
-            )
-            # --- SPEED OPTIMIZATION: Reduce default steps ---
-            num_inference_steps = gr.Slider(
-                label="Inference Steps",
-                info="Fewer steps = faster, more steps = potentially higher quality. Try 10-15 for speed.",
-                minimum=4, maximum=50, step=1, value=15 # Defaulting to 15 instead of 25
-            )
-            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
-            )
+            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)
             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, show_label=True, position=0.5)
     output_seed = gr.Textbox(label="Seed Used", interactive=False, visible=True, scale=1)
     api_base64_output = gr.Textbox(label="API Base64 Output", interactive=False, visible=False)
 
-    # --- 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_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, steps 15
             examples=[ [path, min(2, INTERNAL_PROCESSING_FACTOR), 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, ],
-            outputs=[result_slider, output_seed], # Base64 output ignored by Examples
-            fn=infer,
-            cache_examples="lazy",
-            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
-    run_button.click(
-        fn=infer,
-        inputs=[
-            seed,
-            randomize_seed,
-            input_im,
-            num_inference_steps,
-            upscale_factor_slider, # Use the slider value here
-            controlnet_conditioning_scale,
-        ],
-        outputs=[result_slider, output_seed, api_base64_output],
-        api_name="upscale" # Keep API name
-    )
-
-# Launch the Gradio app
-# Consider increasing queue timeout if compilation adds significant startup time
+            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")
+
 demo.queue(max_size=10).launch(share=False, show_api=True)