|
import spaces |
|
import gradio as gr |
|
import os |
|
import torch |
|
import uuid |
|
|
|
from PIL import Image |
|
from enhance_utils import enhance_image |
|
|
|
|
|
DEFAULT_SRC_PROMPT = "a woman" |
|
DEFAULT_EDIT_PROMPT = "a woman, with red lips, 8k, high quality" |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
def create_demo() -> gr.Blocks: |
|
from inversion_run_adapter import run as adapter_run |
|
|
|
@spaces.GPU(duration=15) |
|
def image_to_image( |
|
input_image_path: str, |
|
input_image_prompt: str, |
|
edit_prompt: str, |
|
seed: int, |
|
w1: float, |
|
num_steps: int, |
|
start_step: int, |
|
guidance_scale: float, |
|
generate_size: int, |
|
lineart_scale: float, |
|
canny_scale: float, |
|
lineart_detect: float, |
|
canny_detect: float, |
|
): |
|
w2 = 1.0 |
|
|
|
input_image = Image.open(input_image_path) |
|
|
|
w2 = 1.0 |
|
run_model = adapter_run |
|
generated_image = run_model( |
|
input_image, |
|
input_image_prompt, |
|
edit_prompt, |
|
generate_size, |
|
seed, |
|
w1, |
|
w2, |
|
num_steps, |
|
start_step, |
|
guidance_scale, |
|
lineart_scale, |
|
canny_scale, |
|
lineart_detect, |
|
canny_detect, |
|
) |
|
enhanced_image = enhance_image(generated_image, False) |
|
|
|
tmpPrefix = "/tmp/gradio/" |
|
|
|
extension = 'png' |
|
|
|
|
|
|
|
|
|
|
|
targetDir = f"{tmpPrefix}output/" |
|
if not os.path.exists(targetDir): |
|
os.makedirs(targetDir) |
|
|
|
enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}" |
|
enhanced_image.save(enhanced_path, quality=100) |
|
|
|
return enhanced_path |
|
|
|
with gr.Blocks() as demo: |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image_path = gr.File(label="Input Image") |
|
with gr.Column(): |
|
generated_image_path = gr.File(label="Download the segment image", interactive=False) |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT) |
|
edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT) |
|
with gr.Accordion("Advanced Options", open=False): |
|
guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale") |
|
seed = gr.Number(label="Seed", value=8) |
|
generate_size = gr.Number(label="Generate Size", value=1024) |
|
lineart_scale = gr.Slider(minimum=0, maximum=5, value=0.8, step=0.1, label="Lineart Weights", visible=True) |
|
canny_scale = gr.Slider(minimum=0, maximum=5, value=0.4, step=0.1, label="Canny Weights", visible=True) |
|
lineart_detect = gr.Number(label="Lineart Detect", value=0.375, visible=True) |
|
canny_detect = gr.Number(label="Canny Detect", value=0.375, visible=True) |
|
with gr.Column(): |
|
num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps") |
|
start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step") |
|
w1 = gr.Number(label="W1", value=2) |
|
g_btn = gr.Button("Edit Image") |
|
|
|
|
|
g_btn.click( |
|
fn=image_to_image, |
|
inputs=[input_image_path, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, lineart_scale, canny_scale, lineart_detect, canny_detect], |
|
outputs=[generated_image_path], |
|
) |
|
|
|
return demo |
