Update app.py

app.py

@@ -10,6 +10,7 @@ from typing import Iterable
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
 from huggingface_hub import hf_hub_download
+from aura_sr import AuraSR
 from gradio_imageslider import ImageSlider
 
 from gradio.themes import Soft
@@ -35,7 +36,7 @@ class OrangeRedTheme(Soft):
         self,
         *,
         primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red, 
+        secondary_hue: colors.Color | str = colors.orange_red, # Use the new color
         neutral_hue: colors.Color | str = colors.slate,
         text_size: sizes.Size | str = sizes.text_lg,
         font: fonts.Font | str | Iterable[fonts.Font | str] = (
@@ -83,6 +84,7 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # --- # Device and CUDA Setup Check ---
@@ -100,17 +102,18 @@ print("Using device:", device)
 MAX_SEED = np.iinfo(np.int32).max
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-# Load Adapters
 pipe.load_lora_weights("prithivMLmods/Kontext-Top-Down-View", weight_name="Kontext-Top-Down-View.safetensors", adapter_name="top-down")
 pipe.load_lora_weights("prithivMLmods/Kontext-Bottom-Up-View", weight_name="Kontext-Bottom-Up-View.safetensors", adapter_name="bottom-up")
 pipe.load_lora_weights("prithivMLmods/Kontext-CAM-Left-View", weight_name="Kontext-CAM-Left-View.safetensors", adapter_name="left-view")
 pipe.load_lora_weights("prithivMLmods/Kontext-CAM-Right-View", weight_name="Kontext-CAM-Right-View.safetensors", adapter_name="right-view")
 pipe.load_lora_weights("starsfriday/Kontext-Remover-General-LoRA", weight_name="kontext_remove.safetensors", adapter_name="kontext-remove")
 
+aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
+
 @spaces.GPU
-def infer(input_image, prompt, lora_adapter, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(input_image, prompt, lora_adapter, upscale_image, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
     """
-    Perform image editing using FLUX Kontext adapters.
+    Perform image editing and optional upscaling, returning a pair for the ImageSlider.
     """
     if not input_image:
         raise gr.Error("Please upload an image for editing.")
@@ -141,6 +144,10 @@ def infer(input_image, prompt, lora_adapter, seed=42, randomize_seed=False, guid
         generator=torch.Generator().manual_seed(seed),
     ).images[0]
 
+    if upscale_image:
+        progress(0.8, desc="Upscaling image...")
+        image = aura_sr.upscale_4x(image)
+
     return (original_image, image), seed, gr.Button(visible=True)
 
 @spaces.GPU
@@ -148,7 +155,7 @@ def infer_example(input_image, prompt, lora_adapter):
     """
     Wrapper function for gr.Examples to call the main infer logic for the slider.
     """
-    (original_image, generated_image), seed, _ = infer(input_image, prompt, lora_adapter)
+    (original_image, generated_image), seed, _ = infer(input_image, prompt, lora_adapter, upscale_image=False)
     return (original_image, generated_image), seed
 
 css="""
@@ -159,7 +166,7 @@ css="""
 #main-title h1 {font-size: 2.1em !important;}
 """
 
-with gr.Blocks() as demo:
+with gr.Blocks(css=css, theme=orange_red_theme) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# **Kontext-Photo-Mate-v2**", elem_id="main-title")
@@ -206,8 +213,7 @@ with gr.Blocks() as demo:
                     )
                     
             with gr.Column():
-                # Use ImageSlider to display (Before, After) tuple correctly
-                output_slider = gr.Image(label="Output Image", interactive=False, format="png")
+                output_slider = ImageSlider(label="Before / After", show_label=False, interactive=False)
                 reuse_button = gr.Button("Reuse this image", visible=False)
 
                 with gr.Row():
@@ -216,6 +222,9 @@ with gr.Blocks() as demo:
                         choices=["Kontext-Top-Down-View", "Kontext-Remover", "Kontext-Bottom-Up-View", "Kontext-CAM-Left-View", "Kontext-CAM-Right-View"],
                         value="Kontext-Top-Down-View"
                     )
+                    
+                with gr.Row():
+                    upscale_checkbox = gr.Checkbox(label="Upscale the final image", value=False)
 
         gr.Examples(
             examples=[
@@ -228,22 +237,21 @@ with gr.Blocks() as demo:
             inputs=[input_image, prompt, lora_adapter],
             outputs=[output_slider, seed],
             fn=infer_example,
-            cache_examples=False,
+            cache_examples="lazy",
             label="Examples"
         )
             
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_image, prompt, lora_adapter, upscale_checkbox, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_slider, seed, reuse_button]
     )
     
-    # Logic to reuse the generated image (2nd image in the tuple)
     reuse_button.click(
         fn=lambda images: images[1] if isinstance(images, (list, tuple)) and len(images) > 1 else images,
         inputs=[output_slider],
         outputs=[input_image]
     )
 
-demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)
+demo.launch(mcp_server=True, ssr_mode=False, show_error=True)