Update app.py

app.py

@@ -9,6 +9,7 @@ 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
 
 # --- Main Model Initialization ---
 MAX_SEED = np.iinfo(np.int32).max
@@ -21,29 +22,22 @@ pipe.load_lora_weights("prithivMLmods/Polaroid-Warm-i2i", weight_name="Polaroid-
 pipe.load_lora_weights("prithivMLmods/Monochrome-Pencil", weight_name="Monochrome-Pencil-i2i.safetensors", adapter_name="pencil")
 
 # --- Upscaler Model Initialization ---
-# FIX: Removed the .to("cuda") call. The library handles device placement automatically.
 aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
 
 @spaces.GPU
 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 and optional upscaling.
-    
-    Args:
-        input_image (PIL.Image.Image): The input image to be edited.
-        prompt (str): Text description of the desired edit.
-        lora_adapter (str): The name of the LoRA adapter to use.
-        upscale_image (bool): If True, the final image will be upscaled 4x.
-        seed (int, optional): Random seed for reproducible generation.
-        randomize_seed (bool, optional): If True, generates a random seed.
-        guidance_scale (float, optional): Controls adherence to the prompt.
-        steps (int, optional): Number of diffusion steps.
-        progress (gr.Progress, optional): Gradio progress tracker.
+    Perform image editing and optional upscaling, returning a pair for the ImageSlider.
     
     Returns:
-        tuple: A 3-tuple containing the generated/upscaled image, the seed used, 
-               and a Gradio update to make the reuse button visible.
+        tuple: A 3-tuple containing:
+               - (PIL.Image.Image, PIL.Image.Image): A tuple of the (original, generated) images for the slider.
+               - int: The seed used for generation.
+               - gr.update: A Gradio update to make the reuse button visible.
     """
+    if not input_image:
+        raise gr.Error("Please upload an image for editing.")
+
     if lora_adapter == "PhotoCleanser":
         pipe.set_adapters(["cleanser"], adapter_weights=[1.0])
     elif lora_adapter == "PhotoRestorer":
@@ -56,39 +50,31 @@ def infer(input_image, prompt, lora_adapter, upscale_image, seed=42, randomize_s
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    if input_image:
-        input_image = input_image.convert("RGB")
-        image = pipe(
-            image=input_image, 
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            width = input_image.size[0],
-            height = input_image.size[1],
-            num_inference_steps=steps,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
-    else:
-        image = pipe(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            generator=torch.Generator().manual_seed(seed),
-        ).images[0]
+    original_image = input_image.copy().convert("RGB")
+    
+    image = pipe(
+        image=original_image, 
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        width = original_image.size[0],
+        height = original_image.size[1],
+        num_inference_steps=steps,
+        generator=torch.Generator().manual_seed(seed),
+    ).images[0]
 
-    # Conditionally upscale the generated image
     if upscale_image:
         progress(0.8, desc="Upscaling image...")
         image = aura_sr.upscale_4x(image)
 
-    return image, seed, gr.Button(visible=True)
+    return (original_image, image), seed, gr.Button(visible=True)
 
 @spaces.GPU
 def infer_example(input_image, prompt, lora_adapter):
     """
-    Wrapper function for gr.Examples. Upscaling is disabled for examples for speed.
+    Wrapper function for gr.Examples to call the main infer logic for the slider.
     """
-    image, seed, _ = infer(input_image, prompt, lora_adapter, upscale_image=False)
-    return image, seed
+    (original_image, generated_image), seed, _ = infer(input_image, prompt, lora_adapter, upscale_image=False)
+    return (original_image, generated_image), seed
 
 css="""
 #col-container {
@@ -145,7 +131,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     )
                     
             with gr.Column():
-                result = gr.Image(label="Result", show_label=False, interactive=False, format="png")
+                # Replace the single image result with the ImageSlider
+                output_slider = ImageSlider(label="Before / After", show_label=False, interactive=False)
                 reuse_button = gr.Button("Reuse this image", visible=False)
                 with gr.Row():
                     lora_adapter = gr.Dropdown(
@@ -156,17 +143,18 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     
         gr.Examples(
             examples=[
-                ["photocleanser/1.png", "[photo content], remove the embroidered pattern from the image while preserving the background and remaining elements, maintaining realism and original details.", "PhotoCleanser"],
-                ["photocleanser/2.png", "[photo content], remove the cat from the image while preserving the background and remaining elements, maintaining realism and original details.", "PhotoCleanser"],
-                ["photorestore/1.png", "[photo content], restore and enhance the image by repairing any damage, scratches, or fading. Colorize the photo naturally while preserving authentic textures and details, maintaining a realistic and historically accurate look.", "PhotoRestorer"],
-                ["photorestore/2.png", "[photo content], restore and enhance the image by repairing any damage, scratches, or fading. Colorize the photo naturally while preserving authentic textures and details, maintaining a realistic and historically accurate look.", "PhotoRestorer"],
-                ["polaroid/1.png", "[photo content], apply a warm, vintage Polaroid-style filter, enhancing the image with nostalgic tones, soft focus, and characteristic light leaks for an authentic, retro feel.", "PolaroidWarm"],
-                ["polaroid/2.png", "[photo content], give the image a classic Polaroid look with warm, saturated colors, gentle fading, and a subtle vignette effect, evoking a sense of timeless memories.", "PolaroidWarm"],
-                ["pencil/1.png", "[photo content], transform the image into a detailed monochrome pencil sketch, emphasizing sharp lines, textures, and shading for a classic hand-drawn look.", "MonochromePencil"],
-                ["pencil/2.png", "[photo content], convert the photo into a realistic graphite pencil drawing, capturing the subject's form and depth with varied strokes and contrast.", "MonochromePencil"]
+                ["photocleanser/1.png", "[photo content], remove the embroidered pattern from the image...", "PhotoCleanser"],
+                ["photocleanser/2.png", "[photo content], remove the cat from the image...", "PhotoCleanser"],
+                ["photorestore/1.png", "[photo content], restore and enhance the image by repairing any damage...", "PhotoRestorer"],
+                ["photorestore/2.png", "[photo content], restore and enhance the image by repairing any damage...", "PhotoRestorer"],
+                ["polaroid/1.png", "[photo content], apply a warm, vintage Polaroid-style filter...", "PolaroidWarm"],
+                ["polaroid/2.png", "[photo content], give the image a classic Polaroid look...", "PolaroidWarm"],
+                ["pencil/1.png", "[photo content], transform the image into a detailed monochrome pencil sketch...", "MonochromePencil"],
+                ["pencil/2.png", "[photo content], convert the photo into a realistic graphite pencil drawing...", "MonochromePencil"]
             ],
             inputs=[input_image, prompt, lora_adapter],
-            outputs=[result, seed],
+            # The output now targets the ImageSlider component
+            outputs=[output_slider, seed],
             fn=infer_example,
             cache_examples=False,
             label="Examples (Image | Prompt | Selected LoRA)"
@@ -176,11 +164,15 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[input_image, prompt, lora_adapter, upscale_checkbox, seed, randomize_seed, guidance_scale, steps],
-        outputs=[result, seed, reuse_button]
+        # The output now targets the ImageSlider component
+        outputs=[output_slider, seed, reuse_button]
     )
+    
+    # Update the reuse function to handle the ImageSlider output
     reuse_button.click(
-        fn=lambda image: image,
-        inputs=[result],
+        # The slider outputs a tuple of images; we want the second one (the generated result)
+        fn=lambda images: images[1] if isinstance(images, (list, tuple)) and len(images) > 1 else images,
+        inputs=[output_slider],
         outputs=[input_image]
     )