File size: 3,125 Bytes
4743900 bd5747f 4743900 bd5747f f2cb7ad bd5747f 4743900 f2cb7ad 4743900 f2cb7ad 4743900 |
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 |
import spaces
import gradio as gr
import torch
import uuid
import os
from PIL import Image
from enhance_utils import enhance_image
DEFAULT_SRC_PROMPT = "a woman, photo"
DEFAULT_EDIT_PROMPT = "a beautiful woman, photo, hollywood style face, 8k, high quality"
device = "cuda" if torch.cuda.is_available() else "cpu"
def create_demo() -> gr.Blocks:
from inversion_run_base import run as base_run
@spaces.GPU(duration=10)
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,
enhance_face: bool = True,
):
w2 = 1.0
input_image = Image.open(input_image_path)
icc_profile = input_image.info.get("icc_profile")
run_model = base_run
res_image = run_model(
input_image,
input_image_prompt,
edit_prompt,
seed,
w1,
w2,
num_steps,
start_step,
guidance_scale,
)
enhanced_image = enhance_image(res_image, enhance_face)
tmpPrefix = "/tmp/gradio/"
extension = 'png'
if enhanced_image.mode == 'RGBA':
extension = 'png'
else:
extension = 'jpg'
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, icc_profile=icc_profile)
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")
enhance_face = gr.Checkbox(label="Enhance Face", value=False)
seed = gr.Number(label="Seed", value=8)
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, enhance_face],
outputs=[generated_image_path],
)
return demo |