app.py

import gradio as gr
import subprocess
import os
import shutil
import sys

target_paths = {
    "data": "/home/user/app/upload/data.zip",
    "data_dir": "/home/user/app/upload/data",
    "config_zip": "/home/user/app/upload/config.zip",
    "configs": "/home/user/app/upload/config",
    "config": "/home/user/app/src/multiview_consist_edit/config/infer_tryon_multi.yaml",
    "output_data": "/home/user/app/image_output_tryon_mvhumannet",
    "output_zip": "/home/user/app/outputs/result.zip",
}

def unzip_data():
    if os.path.exists(target_paths["data"]):
        if os.path.exists(target_paths["data_dir"]):
            shutil.rmtree(target_paths["data_dir"])
        os.makedirs(target_paths["data_dir"], exist_ok=True)
        shutil.unpack_archive(target_paths["data"], target_paths["data_dir"])
        # return target_paths["data_dir"]
    else:
        raise FileNotFoundError("Data file not found at " + target_paths["data"])
    if os.path.exists(target_paths["config_zip"]):
        if os.path.exists(target_paths["configs"]):
            shutil.rmtree(target_paths["configs"])
        os.makedirs(target_paths["configs"], exist_ok=True)
        shutil.unpack_archive(target_paths["config_zip"], target_paths["configs"])
        # return target_paths["configs"]
    else:
        raise FileNotFoundError("Config file not found at " + target_paths["config_zip"])

def zip_outputs():
    if os.path.exists(target_paths["output_zip"]):
        os.remove(target_paths["output_zip"])
    shutil.make_archive(target_paths["output_zip"].replace(".zip", ""), 'zip', root_dir=target_paths["output_data"])
    return target_paths["output_zip"]


def start_inference_stream(config_count):
    config_dir = target_paths["configs"]
    config_files = sorted([
        f for f in os.listdir(config_dir)
        if f.endswith(".yaml") or f.endswith(".yml")
    ])
    if config_count < len(config_files):
        config_files = config_files[:config_count]

    for cfg in config_files:
        src_path = os.path.join(config_dir, cfg)
        shutil.copy(src_path, target_paths["config"])

        process = subprocess.Popen(
            ["python", "src/multiview_consist_edit/infer_tryon_multi.py"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            universal_newlines=True
        )
        
        output = []
        for line in process.stdout:
            output.append(line)
            yield "".join(output)

def install_package(package_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", package_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        output = result.stdout + "\n" + result.stderr
        return output
    except Exception as e:
        return f"Error: {str(e)}"
    

def show_package(pkg_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "show", pkg_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return result.stdout if result.stdout else result.stderr
    except Exception as e:
        return str(e)
    

def uninstall_package(package_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "uninstall", package_name, "-y"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        output = result.stdout + "\n" + result.stderr
        return output
    except Exception as e:
        return f"Error: {str(e)}"


uninstall_package("datasets")
install_package("huggingface_hub==0.25.1")
install_package("diffusers==0.25.1")
print("package version set complete")

def save_files(data_file, config_file):
    os.makedirs(os.path.dirname(target_paths["data"]), exist_ok=True)
    os.makedirs(os.path.dirname(target_paths["config_zip"]), exist_ok=True)

    shutil.copy(data_file.name, target_paths["data"])
    shutil.copy(config_file.name, target_paths["config_zip"])
    unzip_data()
    return "檔案已成功上傳、儲存並解壓縮了！"


with gr.Blocks(theme=gr.themes.Origin()) as demo:
    gr.Markdown("## 請先上傳檔案")
    with gr.Row():
        data_input = gr.File(label="上傳資料壓縮檔", file_types=[".zip"])
        with gr.Column():
            config_input = gr.File(label="Config 壓縮檔", file_types=[".zip"])
            config_count = gr.Number(label="Config 總數", precision=0, value=1)

    upload_button = gr.Button("上傳並儲存")
    output = gr.Textbox(label="狀態")


    gr.Markdown("## Inference")
    with gr.Column():
        log_output = gr.Textbox(label="Inference Log", lines=20)
        infer_btn = gr.Button("Start Inference")

    # gr.Markdown("## Pip Installer")
    # with gr.Column():
    #     with gr.Row():
    #         pkg_input = gr.Textbox(lines=1, placeholder="輸入想安裝的套件名稱，例如 diffusers 或 numpy==1.2.0")
    #         install_output = gr.Textbox(label="Install Output", lines=10)
    #     install_btn = gr.Button("Install Package")
    
    # gr.Markdown("## Pip Uninstaller")
    # with gr.Column():
    #     with gr.Row():
    #         pkg_input2 = gr.Textbox(lines=1, placeholder="輸入想解除安裝的套件名稱，例如 diffusers 或 numpy")
    #         uninstall_output = gr.Textbox(label="Uninstall Output", lines=10)
    #     uninstall_btn = gr.Button("Uninstall Package")
    
    # gr.Markdown("## Pip show")
    # with gr.Column():
    #     with gr.Row():
    #         show_input = gr.Textbox(label="輸入套件名稱（如 diffusers）")
    #         show_output = gr.Textbox(label="套件資訊", lines=10)
    #     show_btn = gr.Button("pip show")
    
    gr.Markdown("## Download results")
    with gr.Column():
        file_output = gr.File(label="點擊下載", interactive=True)
        download_btn = gr.Button("下載結果")

    # show_btn.click(fn=show_package, inputs=show_input, outputs=show_output)
    download_btn.click(fn=zip_outputs, outputs=file_output)
    # install_btn.click(fn=install_package, inputs=pkg_input, outputs=install_output)
    infer_btn.click(fn=start_inference_stream,inputs=config_count, outputs=log_output)
    # uninstall_btn.click(fn=uninstall_package, inputs=pkg_input2, outputs=uninstall_output)
    upload_button.click(fn=save_files,inputs=[data_input, config_input],outputs=output)
demo.launch()