Spaces:
Paused
Paused
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| from torchvision import transforms | |
| import gradio as gr | |
| from src.image_prep import canny_from_pil | |
| from src.pix2pix_turbo import Pix2Pix_Turbo | |
| model = Pix2Pix_Turbo("edge_to_image") | |
| def process(input_image, prompt, low_threshold, high_threshold): | |
| # resize to be a multiple of 8 | |
| new_width = input_image.width - input_image.width % 8 | |
| new_height = input_image.height - input_image.height % 8 | |
| input_image = input_image.resize((new_width, new_height)) | |
| canny = canny_from_pil(input_image, low_threshold, high_threshold) | |
| with torch.no_grad(): | |
| c_t = transforms.ToTensor()(canny).unsqueeze(0).cuda() | |
| output_image = model(c_t, prompt) | |
| output_pil = transforms.ToPILImage()(output_image[0].cpu() * 0.5 + 0.5) | |
| # flippy canny values, map all 0s to 1s and 1s to 0s | |
| canny_viz = 1 - (np.array(canny) / 255) | |
| canny_viz = Image.fromarray((canny_viz * 255).astype(np.uint8)) | |
| return canny_viz, output_pil | |
| if __name__ == "__main__": | |
| # load the model | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Pix2pix-Turbo: **Canny Edge -> Image**") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(sources="upload", type="pil") | |
| prompt = gr.Textbox(label="Prompt") | |
| low_threshold = gr.Slider( | |
| label="Canny low threshold", | |
| minimum=1, | |
| maximum=255, | |
| value=100, | |
| step=10, | |
| ) | |
| high_threshold = gr.Slider( | |
| label="Canny high threshold", | |
| minimum=1, | |
| maximum=255, | |
| value=200, | |
| step=10, | |
| ) | |
| run_button = gr.Button(value="Run") | |
| with gr.Column(): | |
| result_canny = gr.Image(type="pil") | |
| with gr.Column(): | |
| result_output = gr.Image(type="pil") | |
| prompt.submit( | |
| fn=process, | |
| inputs=[input_image, prompt, low_threshold, high_threshold], | |
| outputs=[result_canny, result_output], | |
| ) | |
| low_threshold.change( | |
| fn=process, | |
| inputs=[input_image, prompt, low_threshold, high_threshold], | |
| outputs=[result_canny, result_output], | |
| ) | |
| high_threshold.change( | |
| fn=process, | |
| inputs=[input_image, prompt, low_threshold, high_threshold], | |
| outputs=[result_canny, result_output], | |
| ) | |
| run_button.click( | |
| fn=process, | |
| inputs=[input_image, prompt, low_threshold, high_threshold], | |
| outputs=[result_canny, result_output], | |
| ) | |
| demo.queue() | |
| demo.launch(debug=True) | |