giga_App.py

import gradio as gr
from gradio_imageslider import ImageSlider
from PIL import Image
import numpy as np
from aura_sr import AuraSR
import torch
import os
import time
import platform
import argparse

# Global variable to control batch processing cancellation.
stop_batch_flag = False

def open_folder():
    open_folder_path = os.path.abspath("outputs")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')

def get_placeholder_image():
    """
    Creates a placeholder image (if not already present) and returns its file path.
    This placeholder is a blank (white) image that will be used for progress updates.
    """
    placeholder_path = "placeholder.png"
    if not os.path.exists(placeholder_path):
        placeholder = Image.new("RGB", (256, 256), (255, 255, 255))
        placeholder.save(placeholder_path)
    return placeholder_path

# Force CPU usage
torch.set_default_tensor_type(torch.FloatTensor)

# Override torch.load to always use CPU
original_load = torch.load
torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))

# Initialize the AuraSR model
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")

# Restore original torch.load
torch.load = original_load

def process_single_image(input_image_path, reduce_seams):
    if input_image_path is None:
        raise gr.Error("Please provide an image to upscale.")

    # Send an initial progress update.
    # Instead of (None, None), we use the placeholder image file paths.
    placeholder = get_placeholder_image()
    yield [(placeholder, placeholder), "Starting upscaling..."]

    # Load the image.
    pil_image = Image.open(input_image_path)

    # Upscale using the chosen method.
    start_time = time.time()
    if reduce_seams:
        print("using reduce seams")
        upscaled_image = aura_sr.upscale_4x_overlapped(pil_image)
    else:
        upscaled_image = aura_sr.upscale_4x(pil_image)
    processing_time = time.time() - start_time
    print(f"Processing time: {processing_time:.2f} seconds")

    # Save the upscaled image.
    output_folder = "outputs"
    os.makedirs(output_folder, exist_ok=True)
    
    input_filename = os.path.basename(input_image_path)
    output_filename = os.path.splitext(input_filename)[0]
    output_path = os.path.join(output_folder, output_filename + ".png")
    
    counter = 1
    while os.path.exists(output_path):
        output_path = os.path.join(output_folder, f"{output_filename}_{counter:04d}.png")
        counter += 1
    
    upscaled_image.save(output_path)

    # Send the final progress update along with the before/after slider images.
    yield [(input_image_path, output_path),
           f"Upscaling complete in {processing_time:.2f} seconds"]

def process_batch(input_folder, output_folder=None, reduce_seams=False):
    global stop_batch_flag
    # Reset the stop flag for each new batch process.
    stop_batch_flag = False

    if not input_folder:
        raise gr.Error("Please provide an input folder path.")

    if not output_folder:
        output_folder = "outputs"

    os.makedirs(output_folder, exist_ok=True)
    input_files = [f for f in os.listdir(input_folder) if f.lower().endswith(
        ('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]
    total_files = len(input_files)
    processed_files = 0
    results = []

    yield [results, "Starting batch processing..."]

    for filename in input_files:
        # Check if the stop flag has been set.
        if stop_batch_flag:
            yield [results, "Batch processing cancelled by user."]
            return

        input_path = os.path.join(input_folder, filename)
        pil_image = Image.open(input_path)

        start_time = time.time()
        if reduce_seams:
            upscaled_image = aura_sr.upscale_4x_overlapped(pil_image)
        else:
            upscaled_image = aura_sr.upscale_4x(pil_image)
        processing_time = time.time() - start_time

        output_filename = os.path.splitext(filename)[0] + ".png"
        output_path = os.path.join(output_folder, output_filename)
        
        counter = 1
        while os.path.exists(output_path):
            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}_{counter:04d}.png")
            counter += 1

        upscaled_image.save(output_path)
        processed_files += 1
        results.append(output_path)
        yield [results, f"Processed {processed_files}/{total_files}: {filename} in {processing_time:.2f} seconds"]

    yield [results, f"Batch processing complete. {processed_files} images processed."]

def stop_batch_process():
    global stop_batch_flag
    stop_batch_flag = True
    return "Stop button clicked. Cancelling batch processing..."

title = """<h1 align="center">AuraSR Giga Upscaler V4 by SECourses - Upscales to 4x</h1>
<p><center>AuraSR: new open source super-resolution upscaler based on GigaGAN. Works perfect on some images and fails on some images so give it a try</center></p>
<p><center>Works very fast and very VRAM friendly</center></p>
<h2 align="center">Latest version on : <a href="https://www.patreon.com/posts/121441873">https://www.patreon.com/posts/121441873</a></h2>
"""

def create_demo():
    with gr.Blocks() as demo:
        gr.HTML(title)
        
        with gr.Tab("Single Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(label="Input Image", type="filepath")
                    reduce_seams = gr.Checkbox(
                        label="Reduce Seam Artifacts",
                        value=True,
                        info="upscale_4x upscales the image in tiles that do not overlap. This can result in seams. Use upscale_4x_overlapped to reduce seams. This will double the time upscaling by taking an additional pass and averaging the results."
                    )
                    process_btn = gr.Button(value="Upscale Image", variant="primary")
                with gr.Column(scale=1):
                    # Use the ImageSlider component for comparing before & after images.
                    # "filepath" type means the component expects image file paths.
                    output_slider = ImageSlider(label="Before / After", type="filepath", slider_color="blue")
                    progress_text = gr.Markdown("Progress messages will appear here.")
                    btn_open_outputs = gr.Button("Open Outputs Folder", variant="primary")
                    btn_open_outputs.click(fn=open_folder)
            
            process_btn.click(
                fn=process_single_image,
                inputs=[input_image, reduce_seams],
                outputs=[output_slider, progress_text]
            )
        
        with gr.Tab("Batch Processing"):
            with gr.Row():
                input_folder = gr.Textbox(label="Input Folder Path")
                output_folder = gr.Textbox(label="Output Folder Path (Optional)")
                reduce_seams_batch = gr.Checkbox(
                    label="Reduce Seam Artifacts",
                    value=True,
                    info="upscale_4x upscales the image in tiles that do not overlap. This can result in seams. Use upscale_4x_overlapped to reduce seams. This will double the time upscaling by taking an additional pass and averaging the results."
                )
            with gr.Row():
                batch_process_btn = gr.Button(value="Process Batch", variant="primary")
                stop_batch_btn = gr.Button(value="Stop Batch Processing", variant="secondary")
            with gr.Column():
                output_gallery_batch = gr.Gallery(label="Processed Images")
                progress_text_batch = gr.Markdown("Progress messages will appear here.")
            
            batch_process_btn.click(
                fn=process_batch,
                inputs=[input_folder, output_folder, reduce_seams_batch],
                outputs=[output_gallery_batch, progress_text_batch]
            )
            stop_batch_btn.click(
                fn=stop_batch_process,
                inputs=[],
                outputs=[progress_text_batch]
            )
    
    return demo

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AuraSR Image Upscaling")
    parser.add_argument("--share", action="store_true", help="Create a publicly shareable link")
    args = parser.parse_args()

    demo = create_demo()
    demo.launch(debug=True, inbrowser=True, share=args.share)