Update app.py

app.py

@@ -2,61 +2,30 @@ import spaces
 import os
 import requests
 import time
-import io
+
 import torch
+
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler, DPMSolverMultistepScheduler
+from diffusers.models import AutoencoderKL
+from diffusers.models.attention_processor import AttnProcessor2_0
+
 from PIL import Image
 import cv2
 import numpy as np
-from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler
-from diffusers.models import AutoencoderKL
-from diffusers.models.attention_processor import AttnProcessor2_0
-from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+
 from RealESRGAN import RealESRGAN
+
+import random
+import math
+from scipy.signal import gaussian
+
 import gradio as gr
-import subprocess
-from tqdm import tqdm
-import shutil
-import uuid
-import json
-import threading
-
-# Constants
+from gradio_imageslider import ImageSlider
+
 USE_TORCH_COMPILE = False
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
-# Ensure CUDA is available
-if not torch.cuda.is_available():
-    raise RuntimeError("CUDA is not available. This script requires a CUDA-capable GPU.")
-
-device = torch.device("cuda")
-print(f"Using device: {device}")
-
-# Replace the global abort_status with an Event
-abort_event = threading.Event()
-
-css = """
-.gradio-container {
-    max-width: 100% !important;
-    padding: 20px !important;
-}
-#component-0 {
-    height: auto !important;
-    overflow: visible !important;
-}
-"""
-
-def abort_job():
-    if abort_event.is_set():
-        return "Job is already being aborted."
-    abort_event.set()
-    return "Aborting job... Processing will stop after the current frame."
-
-def check_ffmpeg():
-    try:
-        subprocess.run(["ffmpeg", "-version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
-        return True
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        return False
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def download_file(url, folder_path, filename):
     if not os.path.exists(folder_path):
@@ -102,42 +71,44 @@ def timer_func(func):
         return result
     return wrapper
 
-class ModelManager:
+def get_scheduler(scheduler_name, config):
+    if scheduler_name == "DDIM":
+        return DDIMScheduler.from_config(config)
+    elif scheduler_name == "DPM++ 3M SDE Karras":
+        return DPMSolverMultistepScheduler.from_config(config, algorithm_type="sde-dpmsolver++", use_karras_sigmas=True)
+    elif scheduler_name == "DPM++ 3M Karras":
+        return DPMSolverMultistepScheduler.from_config(config, algorithm_type="dpmsolver++", use_karras_sigmas=True)
+    else:
+        raise ValueError(f"Unknown scheduler: {scheduler_name}")
+
+class LazyLoadPipeline:
     def __init__(self):
         self.pipe = None
-        self.realesrgan_x2 = None
-        self.realesrgan_x4 = None
 
-    def load_models(self):
-        if self.pipe is None:           
+    @timer_func
+    def load(self):
+        if self.pipe is None:
+            print("Starting to load the pipeline...")
             self.pipe = self.setup_pipeline()
+            print(f"Moving pipeline to device: {device}")
             self.pipe.to(device)
-            self.pipe.unet.set_attn_processor(AttnProcessor2_0())
-            self.pipe.vae.set_attn_processor(AttnProcessor2_0())
             if USE_TORCH_COMPILE:
+                print("Compiling the model...")
                 self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
 
-        if self.realesrgan_x2 is None:           
-            self.realesrgan_x2 = RealESRGAN(device, scale=2)
-            self.realesrgan_x2.load_weights('models/upscalers/RealESRGAN_x2.pth', download=False)
-
-        if self.realesrgan_x4 is None:            
-            self.realesrgan_x4 = RealESRGAN(device, scale=4)
-            self.realesrgan_x4.load_weights('models/upscalers/RealESRGAN_x4.pth', download=False)
-
-
+    @timer_func
     def setup_pipeline(self):
+        print("Setting up the pipeline...")
         controlnet = ControlNetModel.from_single_file(
             "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
         )
-        safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
         model_path = "models/models/Stable-diffusion/juggernaut_reborn.safetensors"
         pipe = StableDiffusionControlNetImg2ImgPipeline.from_single_file(
             model_path,
             controlnet=controlnet,
             torch_dtype=torch.float16,
             use_safetensors=True,
-            safety_checker=safety_checker
+            safety_checker=None
         )
         vae = AutoencoderKL.from_single_file(
             "models/VAE/vae-ft-mse-840000-ema-pruned.safetensors",
@@ -154,291 +125,227 @@ class ModelManager:
         pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)
         return pipe
 
-    @timer_func
-    def process_image_batch(self, input_images, resolution, num_inference_steps, strength, hdr, guidance_scale):
-        condition_images = [self.prepare_image(img, resolution, hdr) for img in input_images]
-
-        prompt = "masterpiece, best quality, highres"
-        negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
-
-        options = {
-            "prompt": [prompt] * len(input_images),
-            "negative_prompt": [negative_prompt] * len(input_images),
-            "image": condition_images,
-            "control_image": condition_images,
-            "width": condition_images[0].size[0],
-            "height": condition_images[0].size[1],
-            "strength": strength,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": guidance_scale,
-            "generator": torch.Generator(device=device).manual_seed(0),
-        }
-
-        print("Running inference on batch...")
-        results = self.pipe(**options).images
-        print("Batch processing completed successfully")
-
-        return results
-
-    def prepare_image(self, input_image, resolution, hdr):
-        condition_image = self.resize_and_upscale(input_image, resolution)
-        condition_image = self.create_hdr_effect(condition_image, hdr)
-        return condition_image
+    def set_scheduler(self, scheduler_name):
+        if self.pipe is not None:
+            self.pipe.scheduler = get_scheduler(scheduler_name, self.pipe.scheduler.config)
 
-    @timer_func
-    def resize_and_upscale(self, input_image, resolution):
-        scale = 2 if resolution <= 2048 else 4
-
-        if isinstance(input_image, str):
-            input_image = Image.open(input_image).convert("RGB")
-        elif isinstance(input_image, io.IOBase):
-            input_image = Image.open(input_image).convert("RGB")
-        elif isinstance(input_image, Image.Image):
-            input_image = input_image.convert("RGB")
-        elif isinstance(input_image, np.ndarray):
-            input_image = Image.fromarray(input_image).convert("RGB")
-        else:
-            raise ValueError(f"Unsupported input type for input_image: {type(input_image)}")
+    def __call__(self, *args, **kwargs):
+        return self.pipe(*args, **kwargs)
 
-        W, H = input_image.size
-        k = float(resolution) / min(H, W)
-        H = int(round(H * k / 64.0)) * 64
-        W = int(round(W * k / 64.0)) * 64
-        img = input_image.resize((W, H), resample=Image.LANCZOS)
+class LazyRealESRGAN:
+    def __init__(self, device, scale):
+        self.device = device
+        self.scale = scale
+        self.model = None
 
-        if scale == 2:
-            img = self.realesrgan_x2.predict(img)
-        else:
-            img = self.realesrgan_x4.predict(img)
+    def load_model(self):
+        if self.model is None:
+            self.model = RealESRGAN(self.device, scale=self.scale)
+            self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False)
+    def predict(self, img):
+        self.load_model()
+        return self.model.predict(img)
 
-        return img
+lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
+lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
 
-    @timer_func
-    def create_hdr_effect(self, original_image, hdr):
-        if hdr == 0:
-            return original_image
-        cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
-        factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
-                   1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
-                   1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
-        images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
-        merge_mertens = cv2.createMergeMertens()
-        hdr_image = merge_mertens.process(images)
-        hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
-        hdr_result = Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
-
-        return hdr_result
-
-model_manager = ModelManager()
-model_manager.load_models()  # Ensure models are loaded
-
-def extract_frames(video_path, output_folder):
-    os.makedirs(output_folder, exist_ok=True)
-    command = [
-        'ffmpeg',
-        '-i', video_path,
-        '-vf', 'fps=30',
-        f'{output_folder}/frame_%06d.png'
-    ]
-    subprocess.run(command, check=True)
-
-def frames_to_video(input_folder, output_path, fps, original_video_path):
-    # First, create the video from frames without audio
-    temp_output_path = output_path + "_temp.mp4"
-    video_command = [
-        'ffmpeg',
-        '-framerate', str(fps),
-        '-i', f'{input_folder}/frame_%06d.png',
-        '-c:v', 'libx264',
-        '-pix_fmt', 'yuv420p',
-        temp_output_path
-    ]
-    subprocess.run(video_command, check=True)
-
-    # Then, copy the audio from the original video and add it to the new video
-    final_command = [
-        'ffmpeg',
-        '-i', temp_output_path,
-        '-i', original_video_path,
-        '-c:v', 'copy',
-        '-c:a', 'aac',
-        '-map', '0:v:0',
-        '-map', '1:a:0?',
-        '-shortest',
-        output_path
-    ]
-    subprocess.run(final_command, check=True)
-
-    # Remove the temporary file
-    os.remove(temp_output_path)
+@timer_func
+def resize_and_upscale(input_image, resolution):
+    scale = 2 if resolution <= 2048 else 4
+    input_image = input_image.convert("RGB")
+    W, H = input_image.size
+    k = float(resolution) / min(H, W)
+    H = int(round(H * k / 64.0)) * 64
+    W = int(round(W * k / 64.0)) * 64
+    img = input_image.resize((W, H), resample=Image.LANCZOS)
+    if scale == 2:
+        img = lazy_realesrgan_x2.predict(img)
+    else:
+        img = lazy_realesrgan_x4.predict(img)
+    return img
 
+@timer_func
+def create_hdr_effect(original_image, hdr):
+    if hdr == 0:
+        return original_image
+    cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
+    factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0 - 0.45 * hdr,
+              1.0 - 0.25 * hdr, 1.0, 1.0 + 0.2 * hdr,
+              1.0 + 0.4 * hdr, 1.0 + 0.6 * hdr, 1.0 + 0.8 * hdr]
+    images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
+    merge_mertens = cv2.createMergeMertens()
+    hdr_image = merge_mertens.process(images)
+    hdr_image_8bit = np.clip(hdr_image * 255, 0, 255).astype('uint8')
+    return Image.fromarray(cv2.cvtColor(hdr_image_8bit, cv2.COLOR_BGR2RGB))
+
+lazy_pipe = LazyLoadPipeline()
+lazy_pipe.load()
 
 @timer_func
-def process_video(input_video, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames=None, frame_interval=1, preserve_frames=False, batch_size=8, progress=gr.Progress()):
-    abort_event.clear()  # Clear the abort flag at the start of a new job
-    print("Starting video processing...")
+def progressive_upscale(input_image, target_resolution, steps=3):
+    current_image = input_image.convert("RGB")
+    current_size = max(current_image.size)
     
-
-    # Create a new job folder
-    job_id = str(uuid.uuid4())
-    job_folder = os.path.join("jobs", job_id)
-    os.makedirs(job_folder, exist_ok=True)
-
-    # Save job config
-    config = {
-    "resolution": resolution,
-    "num_inference_steps": num_inference_steps,
-    "strength": strength,
-    "hdr": hdr,
-    "guidance_scale": guidance_scale,
-    "max_frames": max_frames,
-    "frame_interval": frame_interval,
-    "preserve_frames": preserve_frames,
-    "batch_size": batch_size
-    }
-    with open(os.path.join(job_folder, "config.json"), "w") as f:
-        json.dump(config, f)
-
-    # If input_video is a file object or has a 'name' attribute, use its name
-    if isinstance(input_video, io.IOBase) or hasattr(input_video, 'name'):
-        input_video = input_video.name
-
-    # Set up folders
-    frames_folder = os.path.join(job_folder, "video_frames")
-    processed_frames_folder = os.path.join(job_folder, "processed_frames")
-    os.makedirs(frames_folder, exist_ok=True)
-    os.makedirs(processed_frames_folder, exist_ok=True)
-
-    # Extract frames
-    progress(0.1, desc="Extracting frames...")
-    extract_frames(input_video, frames_folder)
-
-    # Process selected frames
-    frame_files = sorted(os.listdir(frames_folder))
-    total_frames = len(frame_files)
-    frames_to_process = min(max_frames, total_frames) if max_frames else total_frames
-
-    try:
-        progress(0.2, desc="Processing frames...")
-        for i in tqdm(range(0, frames_to_process, batch_size), desc="Processing batches"):
-            if abort_event.is_set():
-                print("Job aborted. Stopping processing of new frames.")
-                break
-
-            batch_frames = frame_files[i:min(i+batch_size, frames_to_process)]
-            input_images = [Image.open(os.path.join(frames_folder, frame)) for frame in batch_frames]
-            
-            processed_images = model_manager.process_image_batch(input_images, resolution, num_inference_steps, strength, hdr, guidance_scale)
-
-            for frame_file, processed_image in zip(batch_frames, processed_images):
-                output_frame_path = os.path.join(processed_frames_folder, frame_file)
-                if not preserve_frames or not os.path.exists(output_frame_path):
-                    processed_image.save(output_frame_path)
-                
-            progress((0.2 + 0.7 * (i + batch_size) / frames_to_process), desc=f"Processed batch {i//batch_size + 1}/{(frames_to_process-1)//batch_size + 1}")
-
-        # Always attempt to reassemble video
-        progress(0.9, desc="Reassembling video...")
-        input_filename = os.path.splitext(os.path.basename(input_video))[0]
-        output_video = os.path.join(job_folder, f"{input_filename}_upscaled.mp4")
-        frames_to_video(processed_frames_folder, output_video, 30, input_video)
-
-        if abort_event.is_set():
-            progress(1.0, desc="Video processing aborted, but partial result saved")
-            print("Video processing aborted, but partial result saved")
+    for _ in range(steps):
+        if current_size >= target_resolution:
+            break
+        
+        scale_factor = min(2, target_resolution / current_size)
+        new_size = (int(current_image.width * scale_factor), int(current_image.height * scale_factor))
+        
+        if scale_factor <= 1.5:
+            current_image = current_image.resize(new_size, Image.LANCZOS)
         else:
-            progress(1.0, desc="Video processing completed successfully")
-            print("Video processing completed successfully")
-
-        return output_video
-
-    except Exception as e:
-        print(f"An error occurred during processing: {str(e)}")
-        progress(1.0, desc=f"Error: {str(e)}")
-        return None
-
-@spaces.GPU(duration=200)
-def gradio_process_media(input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size, progress=gr.Progress()):
-    abort_event.clear()  # Clear the abort flag at the start of a new job
-    if input_media is None:
-        return None, "No input media provided."
-
-    print(f"Input media type: {type(input_media)}")
-
-    # Get the file path
-    if isinstance(input_media, str):
-        file_path = input_media
-    elif isinstance(input_media, io.IOBase):
-        file_path = input_media.name
-    elif hasattr(input_media, 'name'):
-        file_path = input_media.name
+            current_image = lazy_realesrgan_x2.predict(current_image)
+        
+        current_size = max(current_image.size)
+    
+    # Final resize to exact target resolution
+    if current_size != target_resolution:
+        aspect_ratio = current_image.width / current_image.height
+        if current_image.width > current_image.height:
+            new_size = (target_resolution, int(target_resolution / aspect_ratio))
+        else:
+            new_size = (int(target_resolution * aspect_ratio), target_resolution)
+        current_image = current_image.resize(new_size, Image.LANCZOS)
+    
+    return current_image
+
+def prepare_image(input_image, resolution, hdr):
+    upscaled_image = progressive_upscale(input_image, resolution)
+    return create_hdr_effect(upscaled_image, hdr)
+
+def create_gaussian_weight(tile_size, sigma=0.3):
+    x = np.linspace(-1, 1, tile_size)
+    y = np.linspace(-1, 1, tile_size)
+    xx, yy = np.meshgrid(x, y)
+    gaussian_weight = np.exp(-(xx**2 + yy**2) / (2 * sigma**2))
+    return gaussian_weight
+
+def adaptive_tile_size(image_size, base_tile_size=512, max_tile_size=1024):
+    w, h = image_size
+    aspect_ratio = w / h
+    if aspect_ratio > 1:
+        tile_w = min(w, max_tile_size)
+        tile_h = min(int(tile_w / aspect_ratio), max_tile_size)
     else:
-        raise ValueError(f"Unsupported input type: {type(input_media)}")
+        tile_h = min(h, max_tile_size)
+        tile_w = min(int(tile_h * aspect_ratio), max_tile_size)
+    return max(tile_w, base_tile_size), max(tile_h, base_tile_size)
 
-    print(f"File path: {file_path}")
+def process_tile(tile, num_inference_steps, strength, guidance_scale, controlnet_strength):
+    prompt = "masterpiece, best quality, highres"
+    negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
+    
+    options = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "image": tile,
+        "control_image": tile,
+        "num_inference_steps": num_inference_steps,
+        "strength": strength,
+        "guidance_scale": guidance_scale,
+        "controlnet_conditioning_scale": float(controlnet_strength),
+        "generator": torch.Generator(device=device).manual_seed(random.randint(0, 2147483647)),
+    }
+    
+    return np.array(lazy_pipe(**options).images[0])
 
-    # Check if the file is a video
-    video_extensions = ('.mp4', '.avi', '.mov', '.mkv')
-    if file_path.lower().endswith(video_extensions):
-        print("Processing video...")
-        result = process_video(file_path, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size, progress)
-        if result:
-            return result, "Video processing completed successfully."
-        else:
-            return None, "Error occurred during video processing."
-    else:
-        print("Processing image...")
-        result = model_manager.process_image(file_path, resolution, num_inference_steps, strength, hdr, guidance_scale)
-        if result:
-            # Save the processed image
-            output_path = os.path.join("processed_images", f"processed_{os.path.basename(file_path)}")
-            os.makedirs(os.path.dirname(output_path), exist_ok=True)
-            result.save(output_path)
-            return output_path, "Image processing completed successfully."
-        else:
-            return None, "Error occurred during image processing."
-
-title = """
-    <h1 align="center">Simple Slow Video Upscaler</h1>
-    <p align="center">Currently takes too long :( Please try 1-2 seconds videos only.</p>
-    <p align="center">
-    <a href="https://twitter.com/hrishioa" target="_blank">[Hrishi]</a>
-    <a href="https://huggingface.co/spaces/gokaygokay/Tile-Upscaler" target="_blank">[gokaygokay/Tile-Upscaler]</a>
-    <a href="https://github.com/philz1337x/clarity-upscaler" target="_blank">[philz1337x]</a>
-    <a href="https://github.com/BatouResearch/controlnet-tile-upscale" target="_blank">[BatouResearch]</a>
-    <a href="https://github.com/hrishioa/SimpleSlowVideoUpscaler" target="_blank">[Hrishi GitHub]</a>
-    </p>
-    """
-# Update the Gradio interface
-with gr.Blocks(css=css, theme=gr.themes.Default(primary_hue="blue")) as iface:
-    gr.HTML(title)
+@spaces.GPU
+@timer_func
+def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale, controlnet_strength, scheduler_name):
+    print("Starting image processing...")
+    torch.cuda.empty_cache()
+    lazy_pipe.set_scheduler(scheduler_name)
+    
+    # Convert input_image to numpy array
+    input_array = np.array(input_image)
+    
+    # Prepare the condition image
+    condition_image = prepare_image(input_image, resolution, hdr)
+    W, H = condition_image.size
+    
+    # Adaptive tiling
+    tile_width, tile_height = adaptive_tile_size((W, H))
+    
+    # Calculate the number of tiles
+    overlap = min(64, tile_width // 8, tile_height // 8)  # Adaptive overlap
+    num_tiles_x = math.ceil((W - overlap) / (tile_width - overlap))
+    num_tiles_y = math.ceil((H - overlap) / (tile_height - overlap))
+    
+    # Create a blank canvas for the result
+    result = np.zeros((H, W, 3), dtype=np.float32)
+    weight_sum = np.zeros((H, W, 1), dtype=np.float32)
+    
+    # Create gaussian weight
+    gaussian_weight = create_gaussian_weight(max(tile_width, tile_height))
+    
+    for i in range(num_tiles_y):
+        for j in range(num_tiles_x):
+            # Calculate tile coordinates
+            left = j * (tile_width - overlap)
+            top = i * (tile_height - overlap)
+            right = min(left + tile_width, W)
+            bottom = min(top + tile_height, H)
+            
+            # Adjust tile size if it's at the edge
+            current_tile_size = (bottom - top, right - left)
+            
+            tile = condition_image.crop((left, top, right, bottom))
+            tile = tile.resize((tile_width, tile_height))
+            
+            # Process the tile
+            result_tile = process_tile(tile, num_inference_steps, strength, guidance_scale, controlnet_strength)
+            
+            # Apply gaussian weighting
+            if current_tile_size != (tile_width, tile_height):
+                result_tile = cv2.resize(result_tile, current_tile_size[::-1])
+                tile_weight = cv2.resize(gaussian_weight, current_tile_size[::-1])
+            else:
+                tile_weight = gaussian_weight[:current_tile_size[0], :current_tile_size[1]]
+            
+            # Add the tile to the result with gaussian weighting
+            result[top:bottom, left:right] += result_tile * tile_weight[:, :, np.newaxis]
+            weight_sum[top:bottom, left:right] += tile_weight[:, :, np.newaxis]
+    
+    # Normalize the result
+    final_result = (result / weight_sum).astype(np.uint8)
+    
+    print("Image processing completed successfully")
+    
+    return [input_array, final_result]
+
+title = """<h1 align="center">Tiled Upscaler V2</h1>
+<p align="center">The main ideas come from</p>
+<p><center>
+<a href="https://github.com/philz1337x/clarity-upscaler" target="_blank">[philz1337x]</a>
+<a href="https://github.com/BatouResearch/controlnet-tile-upscale" target="_blank">[Pau-Lozano]</a>
+</center></p>
+"""
 
+with gr.Blocks() as demo:
+    gr.HTML(title)
     with gr.Row():
-        with gr.Column(scale=2):
-            input_media = gr.File(label="Input Media (Image or Video)")
-            resolution = gr.Slider(256, 2048, 512, step=256, label="Resolution")
-            num_inference_steps = gr.Slider(1, 50, 20, step=1, label="Number of Inference Steps")
-            strength = gr.Slider(0, 1, 0.05, step=0.01, label="Strength")
-            hdr = gr.Slider(0, 1, 0, step=0.1, label="HDR Effect")
-            guidance_scale = gr.Slider(0, 20, 5, step=0.5, label="Guidance Scale")
-            max_frames = gr.Number(label="Max Frames to Process (leave empty for full video)", precision=0)
-            frame_interval = gr.Slider(1, 30, 1, step=1, label="Frame Interval (process every nth frame)")
-            preserve_frames = gr.Checkbox(label="Preserve Existing Processed Frames", value=True)
-            batch_size = gr.Slider(1, 16, 8, step=1, label="Batch Size")
-
-        with gr.Column(scale=1):
-            submit_button = gr.Button("Process Media")
-            abort_button = gr.Button("Abort Job")
-            output = gr.File(label="Processed Media")
-            status = gr.Markdown("Ready to process media.")
-
-    submit_button.click(
-        gradio_process_media,
-        inputs=[input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size],
-        outputs=[output, status]
-    )
-
-    abort_button.click(abort_job, inputs=[], outputs=status)
-
-# Launch the Gradio app
-iface.launch()
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            run_button = gr.Button("Enhance Image")
+        with gr.Column():
+            output_slider = ImageSlider(label="Before / After", type="numpy")
+    with gr.Accordion("Advanced Options", open=False):
+        resolution = gr.Slider(minimum=128, maximum=2048, value=1536, step=128, label="Resolution")
+        num_inference_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Inference Steps")
+        strength = gr.Slider(minimum=0, maximum=1, value=0.4, step=0.01, label="Strength")
+        hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect")
+        guidance_scale = gr.Slider(minimum=0, maximum=20, value=6, step=0.5, label="Guidance Scale")
+        controlnet_strength = gr.Slider(minimum=0.0, maximum=2.0, value=0.75, step=0.05, label="ControlNet Strength")
+        scheduler_name = gr.Dropdown(
+            choices=["DDIM", "DPM++ 3M SDE Karras", "DPM++ 3M Karras"],
+            value="DDIM",
+            label="Scheduler"
+        )
+
+    run_button.click(fn=gradio_process_image, 
+                     inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale, controlnet_strength, scheduler_name],
+                     outputs=output_slider)
+
+demo.launch(debug=True, share=True)