Update app.py

app.py

@@ -2,8 +2,8 @@ import logging
 import random
 import warnings
 import os
-import io
-import base64
+import io # Ensure io is imported
+import base64 # Ensure base64 is imported
 import gradio as gr
 import numpy as np
 import spaces
@@ -21,10 +21,10 @@ warnings.filterwarnings("ignore")
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 512px; /* Increased max-width slightly for better layout */
+    max-width: 512px;
 }
 .gradio-container {
-    max-width: 900px !important; /* Control overall container width */
+    max-width: 900px !important;
     margin: auto !important;
 }
 """
@@ -75,20 +75,67 @@ 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"):
+        logging.info("Attempting to compile the pipeline transformer with torch.compile...")
+        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)
+            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}")
+        except Exception as e:
+            logging.warning(f"torch.compile failed: {e}. Running unoptimized.")
+    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}.")
+
+    logging.info("Pipeline ready for inference.")
+
+
 except Exception as 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}")
+    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
-# --- NEW: Define the internal factor for quality ---
+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.
 INTERNAL_PROCESSING_FACTOR = 4
 
-# --- Image Processing Function (Modified) ---
+# --- Image Processing Function (Uses INTERNAL_PROCESSING_FACTOR for budgeting) ---
 def process_input(input_image):
     """Processes the input image for the pipeline.
        The pixel budget check uses the fixed INTERNAL_PROCESSING_FACTOR."""
@@ -126,9 +173,8 @@ def process_input(input_image):
             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)
-            # Ensure minimum size of 8x8
-            input_w_resized = max(8, input_w_resized)
-            input_h_resized = max(8, input_h_resized)
+            input_w_resized = max(8, input_w_resized) # Ensure min size
+            input_h_resized = max(8, input_h_resized) # Ensure min size
             intermediate_w = input_w_resized * INTERNAL_PROCESSING_FACTOR
             intermediate_h = input_h_resized * INTERNAL_PROCESSING_FACTOR
 
@@ -154,14 +200,14 @@ def process_input(input_image):
 
     return input_image_to_process, w_original, h_original, was_resized
 
-# --- Inference Function (Modified) ---
-@spaces.GPU(duration=90)
+# --- Inference Function (Runs pipeline at Internal Factor, resizes to Final Factor) ---
+@spaces.GPU(duration=70) # Keep GPU decorator
 def infer(
     seed,
     randomize_seed,
     input_image,
-    num_inference_steps,
-    final_upscale_factor, # Renamed for clarity internally
+    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,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -170,7 +216,7 @@ def infer(
         gr.Error("Pipeline not loaded. Cannot perform inference.")
         return [[None, None], 0, None]
 
-    original_input_pil = input_image # Keep ref for slider
+    original_input_pil = input_image
 
     if input_image is None:
          gr.Warning("Please provide an input image.")
@@ -180,12 +226,15 @@ def infer(
         seed = random.randint(0, MAX_SEED)
     seed = int(seed)
 
-    # Ensure final_upscale_factor is an integer
+    # Ensure factors are integers
     final_upscale_factor = int(final_upscale_factor)
+    num_inference_steps = int(num_inference_steps) # Ensure steps are int
+
+    # Sanity check: final factor shouldn't exceed internal processing factor in this workflow
     if 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
+        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.")
+        final_upscale_factor = INTERNAL_PROCESSING_FACTOR
 
     logging.info(
         f"Starting inference with seed: {seed}, "
@@ -195,7 +244,6 @@ def infer(
     )
 
     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
         )
@@ -210,20 +258,18 @@ def infer(
     control_image_w = w_proc * INTERNAL_PROCESSING_FACTOR
     control_image_h = h_proc * INTERNAL_PROCESSING_FACTOR
 
-    # 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
+    # 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
         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 to {control_image_w}x{control_image_h} post-processing to fit budget.")
+        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}.")
 
     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}")
@@ -233,26 +279,28 @@ def infer(
     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})...")
-    logging.info(f"Running pipeline with size: {control_image_w}x{control_image_h}")
+    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
     try:
         with torch.inference_mode():
             intermediate_result_image = pipe(
                 prompt="",
-                control_image=control_image, # Control image IS the intermediate size
+                control_image=control_image,
                 controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-                num_inference_steps=int(num_inference_steps),
+                num_inference_steps=num_inference_steps, # Use the integer value
                 guidance_scale=0.0,
-                height=control_image_h, # Target height for the model
-                width=control_image_w,   # Target width for the model
+                height=control_image_h,
+                width=control_image_w,
                 generator=generator,
+                # Add progress callback if desired, requires tqdm
+                # callback_on_step_end = ...
             ).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:
          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}.")
+         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()
          return [[original_input_pil, None], seed, None]
     except Exception as e:
@@ -266,61 +314,52 @@ def infer(
          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
         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}.")
     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
 
     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:
             if final_target_w > 0 and final_target_h > 0:
-                 # 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:
                  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
+                 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}).")
-            final_result_image = intermediate_result_image # Fallback to intermediate
+            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(f"Inference successful. Final output size: {final_result_image.size}")
 
-    # --- Base64 Encoding (No change needed here) ---
+    # --- Base64 Encoding ---
     base64_string = None
     if final_result_image:
         try:
             buffered = io.BytesIO()
-            final_result_image.save(buffered, format="WEBP", quality=90)
+            final_result_image.save(buffered, format="WEBP", quality=90) # WEBP is usually smaller
             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).")
         except Exception as enc_err:
              logging.error(f"Failed to encode result image to Base64: {enc_err}", exc_info=True)
 
-    # Return original input and the FINAL processed image
     return [[original_input_pil, final_result_image], seed, base64_string]
 
 
-# --- Gradio Interface (Minor Text Updates) ---
+# --- Gradio Interface ---
 with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as demo:
     gr.Markdown(
         f"""
@@ -328,11 +367,16 @@ 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,
+    **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, not the final output size.
+    The *diffusion process time* is mainly determined by this intermediate size and the number of steps.
     """
     )
 
@@ -346,9 +390,25 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
             )
         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, 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)
+            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
+            )
             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)
@@ -373,14 +433,14 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
 
     if example_paths:
         gr.Examples(
-            # 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 ],
+            # 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="Example Images (Click to Run with 2x Output)",
+            label=f"Example Images (Click to Run with {min(2, INTERNAL_PROCESSING_FACTOR)}x Output, 15 Steps)",
             run_on_click=True
         )
     else:
@@ -401,8 +461,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Flux Upscaler Demo") as d
             controlnet_conditioning_scale,
         ],
         outputs=[result_slider, output_seed, api_base64_output],
-        api_name="upscale"
+        api_name="upscale" # Keep API name
     )
 
 # Launch the Gradio app
+# Consider increasing queue timeout if compilation adds significant startup time
 demo.queue(max_size=10).launch(share=False, show_api=True)