"vscode:/vscode.git/clone" did not exist on "a4489ffb8ecd547e5e44f22284444348916be012"
Commit 83222671 authored by chenpangpang's avatar chenpangpang
Browse files

feat: zancun

parent 43f12e23
import gradio as gr import gradio as gr
from PIL import Image from PIL import Image
import numpy as np
from aura_sr import AuraSR from aura_sr import AuraSR
import torch
import os
# # Force CPU usage
# torch.set_default_type(torch.FloatTensor)
# torch.set_default_device('cpu')
#
# # 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 USE_TORCH_COMPILE = False
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2/model.safetensors") ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# # Restore original torch.load # Initialize the AuraSR model
# torch.load = original_load aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2/model.safetensors")
def process_image(input_image): def process_image(input_image):
if input_image is None: if input_image is None:
raise gr.Error("Please provide an image to upscale.") raise gr.Error("Please provide an image to upscale.")
print("get input image: ", input_image)
# Convert to PIL Image for resizing
pil_image = Image.fromarray(input_image)
# Upscale the image using AuraSR # Upscale the image using AuraSR
upscaled_image = process_image_on_gpu(pil_image) upscaled_image = process_image_on_gpu(input_image)
print("upscaled_image: ", upscaled_image)
return upscaled_image return upscaled_image
def process_image_on_gpu(pil_image): def process_image_on_gpu(pil_image):
return aura_sr.upscale_4x(pil_image) return aura_sr.upscale_4x(pil_image)
title = """<h1 align="center">AuraSR-v2:一款基于GAN图像修复工具,可从低分辨率图片生成高分辨率图片</h1>""" title = """<h1 align="center">AuraSR-v2:一款基于GAN图像修复工具,可从低分辨率图片生成高分辨率图片</h1>"""
with gr.Blocks() as demo: with gr.Blocks() as demo:
gr.HTML(title) gr.HTML(title)
with gr.Row(): with gr.Row():
with gr.Column(scale=1): with gr.Column(scale=1):
input_image = gr.Image(label="输入图片", type="numpy") input_image = gr.Image(type="pil", label="输入图片")
process_btn = gr.Button("生成") process_btn = gr.Button("生成")
with gr.Column(scale=1): with gr.Column(scale=1):
gallery = gr.Image(label="生成图片") gallery = gr.Image(label="生成图片")
...@@ -61,7 +54,7 @@ with gr.Blocks() as demo: ...@@ -61,7 +54,7 @@ with gr.Blocks() as demo:
inputs=input_image, inputs=input_image,
outputs=gallery, outputs=gallery,
fn=process_image, fn=process_image,
cache_examples=True cache_examples=True,
) )
demo.launch(server_name='0.0.0.0', share=True) demo.launch(server_name='0.0.0.0', share=True)
spaces
aura-sr aura-sr
diffusers
gradio==3.40.0
\ No newline at end of file
...@@ -18,7 +18,7 @@ ...@@ -18,7 +18,7 @@
"\n", "\n",
"## 功能介绍\n", "## 功能介绍\n",
"- 原项目地址:https://huggingface.co/spaces/gokaygokay/AuraSR-v2\n", "- 原项目地址:https://huggingface.co/spaces/gokaygokay/AuraSR-v2\n",
"- AuraSR-v2:一款基于GAN图像修复工具,可从低分辨率图片生成高分辨率图片" "- AuraSR-v2:一款基于GAN图像修复工具,可从低分辨率图片生成高分辨率图片"
] ]
}, },
{ {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment