File size: 2,525 Bytes
5ca78e3 6fa6cf8 bbbd3d2 5ca78e3 12f1b5a 5ca78e3 bbbd3d2 5ca78e3 12f1b5a 5ca78e3 d69b2fc 5ca78e3 d69b2fc 5ca78e3 2d74701 5ca78e3 12f1b5a 5ca78e3 2d74701 5ca78e3 12f1b5a 2d74701 5ca78e3 fdded3f |
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 |
import gradio as gr
import os
import tempfile
import zipfile
import shutil
from test_api import run_adain
def process_adain_api(files, style_strength=1.0, dataset_size=100):
'''API endpoint for AdaIN processing.'''
if not files:
return None, "No files uploaded"
# Create temporary directories
temp_dir = tempfile.mkdtemp()
input_dir = os.path.join(temp_dir, "input")
output_dir = os.path.join(temp_dir, "output")
os.makedirs(input_dir, exist_ok=True)
os.makedirs(output_dir, exist_ok=True)
# Save uploaded files
for file in files:
if file is not None:
shutil.copy(file.name, input_dir)
try:
style_dataset_pth = "tidalove/paleo-real"
# Run AdaIN
run_adain(input_dir, style_dataset_pth, output_dir, style_strength, dataset_size)
# Create zip file
zip_path = os.path.join(temp_dir, "style_transfer_results.zip")
with zipfile.ZipFile(zip_path, 'w') as zipf:
for file in os.listdir(output_dir):
if file.lower().endswith(('.jpg', '.jpeg', '.png')):
zipf.write(os.path.join(output_dir, file), file)
return [os.path.join(output_dir, basename) for basename in os.listdir(output_dir)], zip_path, f"Style transfer completed with strength {style_strength}"
except Exception as e:
return None, f"Error: {str(e)}"
# Create interface with API endpoint
with gr.Blocks() as adain_demo:
gr.Markdown("# AdaIN Style Transfer Service")
with gr.Row():
with gr.Column():
files_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
strength_input = gr.Slider(0.0, 2.0, 1.0, step=0.1, label="Style strength")
size_input = gr.Slider(1, 1000, 100, step=10, label="Generated dataset size")
process_btn = gr.Button("Process", variant="primary")
with gr.Column():
gallery_output = gr.Gallery(label="Gallery", columns=3, rows=2, object_fit="contain")
download_output = gr.File(label="Download Results")
status_output = gr.Textbox(label="Status", interactive=False)
process_btn.click(
fn=process_adain_api,
inputs=[files_input, strength_input, size_input],
outputs=[gallery_output, download_output, status_output],
api_name="adain_process" # This creates the API endpoint
)
adain_demo.launch(show_error=True) |