Generative-AI / giga_App.py
MonsterMMORPG's picture
Upload 2 files
08bd837 verified
raw
history blame
4.9 kB
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)