Spaces:
Runtime error
Runtime error
File size: 3,091 Bytes
d858831 4479682 0022789 9148f31 44f737a 0022789 b2cd494 ff46a0e 44f737a 9148f31 b2cd494 d858831 0022789 813561e ff46a0e 82192ca 0022789 813561e 0022789 ff46a0e 0022789 d858831 813561e adcd556 f1841f9 154494a f1841f9 154494a f1841f9 154494a f1841f9 0022789 d858831 f1841f9 0022789 d858831 0022789 d858831 0022789 d858831 0022789 c5349e7 92f8a2c 44f737a c5349e7 44f737a |
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 |
import os
import time
import spaces
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 devicetorch
DEVICE = devicetorch.get(torch)
# 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", device=DEVICE)
# Restore original torch.load
torch.load = original_load
# Create output folder if not exists
output_folder = '../outputs'
os.makedirs(output_folder, exist_ok=True)
def generate_output_filename():
timestamp = time.strftime("%Y%m%d-%H%M%S")
return f"upscaled_{timestamp}.png"
def process_image(input_image):
if input_image is None:
raise gr.Error("Please provide an image to upscale.")
# Convert to PIL Image for resizing
pil_image = Image.fromarray(input_image)
# Upscale the image using AuraSR
upscaled_image = process_image_on_gpu(pil_image)
# Convert result to numpy array if it's not already
result_array = np.array(upscaled_image)
# Save result as PNG
output_filename = generate_output_filename()
output_path = os.path.join(output_folder, output_filename)
upscaled_image.save(output_path, format="PNG")
return [input_image, result_array], output_path
@spaces.GPU
def process_image_on_gpu(pil_image):
return aura_sr.upscale_4x(pil_image)
title = """<h1 align="center">AuraSR-v2 - An open reproduction of the GigaGAN Upscaler from fal.ai</h1>
<p><center>
<a href="https://huggingface.co/fal/AuraSR-v2" target="_blank">[AuraSR-v2]</a>
<a href="https://blog.fal.ai/introducing-aurasr-an-open-reproduction-of-the-gigagan-upscaler-2/" target="_blank">[Blog Post]</a>
<a href="https://huggingface.co/fal-ai/AuraSR" target="_blank">[v1 Model Page]</a>
</center></p>
"""
css = """
.img { max-height: 80vh !Important }
#slider-container { overflow: hidden; display: flex; justify-content: center }
#row-height { height: 65px !important }
"""
with gr.Blocks(css=css) as demo:
gr.HTML(title)
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(label="Input Image", type="numpy", elem_classes="img")
process_btn = gr.Button("Upscale Image")
with gr.Column(scale=1):
output_slider = ImageSlider(label="Before / After", type="numpy", elem_classes="img", elem_id="slider-container")
download_btn = gr.File(label="Download Image", elem_id="row-height")
process_btn.click(
fn=process_image,
inputs=[input_image],
outputs=[output_slider, download_btn]
)
# Add examples
gr.Examples(
examples=[
"image1.png",
"image3.png"
],
inputs=input_image,
outputs=output_slider,
fn=process_image,
# cache_examples=True
)
demo.launch(debug=True)
|