File size: 4,895 Bytes
08bd837 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import gradio as gr
from PIL import Image
import numpy as np
from aura_sr import AuraSR
import torch
import os
import time
from pathlib import Path
import argparse
# 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):
if input_image_path is None:
raise gr.Error("Please provide an image to upscale.")
# Load the image
pil_image = Image.open(input_image_path)
# Upscale the image using AuraSR
start_time = time.time()
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)
return [input_image_path, output_path]
def process_batch(input_folder, output_folder=None):
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 = []
for filename in input_files:
input_path = os.path.join(input_folder, filename)
pil_image = Image.open(input_path)
start_time = time.time()
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
print(f"Processed {processed_files}/{total_files}: {filename} in {processing_time:.2f} seconds")
results.append(output_path)
print(f"Batch processing complete. {processed_files} images processed.")
return results
title = """<h1 align="center">AuraSR Giga Upscaler V1 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/110060645">https://www.patreon.com/posts/110060645</a></h1>
"""
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")
process_btn = gr.Button(value="Upscale Image", variant="primary")
with gr.Column(scale=1):
output_gallery = gr.Gallery(label="Before / After", columns=2)
process_btn.click(
fn=process_single_image,
inputs=[input_image],
outputs=output_gallery
)
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)")
batch_process_btn = gr.Button(value="Process Batch", variant="primary")
output_gallery = gr.Gallery(label="Processed Images")
batch_process_btn.click(
fn=process_batch,
inputs=[input_folder, output_folder],
outputs=output_gallery
)
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) |