app.py

import gradio as gr
import os, requests
import numpy as np
from inference import setup_model, colorize_grayscale, predict_anchors

## local |  remote
RUN_MODE = "remote"
if RUN_MODE != "local":
    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")
    os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")
    ## examples
    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/04.jpg")

## step 1: set up model
device = "cpu"
checkpt_path = "checkpoints/disco-beta.pth.rar"
colorizer, colorLabeler = setup_model(checkpt_path, device=device)

def click_colorize(rgb_img, hint_img, n_anchors, is_high_res, is_editable):
    if hint_img is None:
        hint_img = rgb_img
    output = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, True, is_editable, device)
    output1 = colorize_grayscale(colorizer, colorLabeler, rgb_img, hint_img, n_anchors, False, is_editable, device)
    return output, output1

def click_predanchors(rgb_img, n_anchors, is_high_res, is_editable):
    output = predict_anchors(colorizer, colorLabeler, rgb_img, n_anchors, is_high_res, is_editable, device)
    return output

## step 2: configure interface
def switch_states(is_checked):
    if is_checked:
        return gr.Image.update(visible=True), gr.Button.update(visible=True)
    else:
        return gr.Image.update(visible=False), gr.Button.update(visible=False)

demo = gr.Interface(title="DISCO")
with demo:
    gr.Markdown(value="""
                    Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors. Check our [project page](https://menghanxia.github.io/projects/disco.html) 😛.
                    """)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                Image_input = gr.inputs.Image(type="numpy", label="Input", interactive=True)
                Image_anchor = gr.inputs.Image(type="numpy", label="Anchor", tool="color-sketch", interactive=True, visible=False)
            with gr.Row():
                Num_anchor = gr.inputs.Number(type="int", value=8, label="Num. of anchors (3~14)")
                Radio_resolution = gr.inputs.Radio(choices=["Low (256x256)", "High (512x512)"], label="Colorization resolution (Low is more stable)", default="Low (256x256)")
            with gr.Row():
                Checkbox_editable = gr.inputs.Checkbox(default=False, label='Show editable anchors')
                Button_show_anchor = gr.inputs.Button(value="Predict anchors", visible=False)
            Button_run = gr.inputs.Button(value="Colorize")

        with gr.Column():
            Image_output = [gr.outputs.Image(type="numpy", label="Output").style(height=480), gr.outputs.Image(type="numpy", label="Output").style(height=480)]

    output_image_anchor, output_button_show_anchor = gr.outputs.Image(), gr.outputs.Button()
    output_image_output = [gr.outputs.Image(label="Output").style(height=480), gr.outputs.Image(label="Output").style(height=480)]

    gr.Interface.Transforms(
        output_image_anchor, output_button_show_anchor, func=switch_states, inputs=Checkbox_editable, outputs=[Image_anchor, Button_show_anchor]
    )

    gr.Interface.Button(Button_show_anchor, func=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Checkbox_editable], outputs=Image_anchor)

    gr.Interface.Button(Button_run, func=click_colorize, inputs=[Image_input, Image_anchor, Num_anchor, Radio_resolution, Checkbox_editable], outputs=output_image_output)

    ## guiline
    if RUN_MODE != "local":
        gr.Interface.Examples(examples=[
            ['04.jpg', 8, "Low (256x256)"],
        ], inputs=[Image_input, Num_anchor, Radio_resolution], outputs=output_image_output, title="Используйте эту фотографию для примера кликнув по ней").initialize(auto_close=False, auto_close_timeout=999)

if RUN_MODE == "local":
    demo.launch(server_name="9.134.253.83", server_port=7788)
else:
    demo.launch()