app.py

import gradio as gr
import os
import torch
import modules.wdtagger

from modules.model import get_model_and_vae_options
from modules.text2img import generate_image_wrapper

# Mendapatkan daftar model dan VAE untuk dropdown
all_models, all_vaes = get_model_and_vae_options()


# Daftar model dan ControlNet
models = ["Model A", "Model B", "Model C"]
vae = ["VAE A", "VAE B", "VAE C"]
controlnet_types = ["Canny", "Depth", "Normal", "Pose"]
schedulers = ["Euler", "LMS", "DDIM"]


# Fungsi placeholder
def load_model(selected_model):
    return f"Model {selected_model} telah dimuat."

def generate_image(prompt, neg_prompt, width, height, scheduler, num_steps, num_images, cfg_scale, seed, model):
    # Logika untuk menghasilkan gambar dari teks menggunakan model
    return [f"Gambar {i+1} untuk prompt '{prompt}' dengan model '{model}'" for i in range(num_images)], {"prompt": prompt, "neg_prompt": neg_prompt}

def process_image(image, prompt, neg_prompt, model):
    # Logika untuk memproses gambar menggunakan model
    return f"Proses gambar dengan prompt '{prompt}' dan model '{model}'"

def controlnet_process(image, controlnet_type, model):
    # Logika untuk memproses gambar menggunakan ControlNet
    return f"Proses gambar dengan ControlNet '{controlnet_type}' dan model '{model}'"


def controlnet_process_func(image, controlnet_type, model):
    # Update fungsi sesuai kebutuhan
    return controlnet_process(image, controlnet_type, model)

def intpaint_func (image, controlnet_type, model):
    # Update fungsi sesuai kebutuhan
    return controlnet_process(image, controlnet_type, model)
       
def intpaint_func (image, controlnet_type, model):
    # Update fungsi sesuai kebutuhan
    return controlnet_process(image, controlnet_type, model)


with gr.Blocks(css="style.css") as app:
    # Dropdown untuk memilih model di luar tab dengan lebar kecil
    with gr.Row():
        checkpoint_dropdown = gr.Dropdown(choices=all_models, label="Model", value=all_models[0])
        vae_dropdown = gr.Dropdown(choices=all_vaes, label="VAE", value=all_vaes[0])

    # Tab untuk Text-to-Image
    with gr.Tab("Text-to-Image"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(label="Prompt", placeholder="Enter Prompt", lines=2, elem_id="prompt-input")
                neg_prompt_input = gr.Textbox(label="Negative prompt", placeholder="Enter Negative Prompt (optional)", lines=2, elem_id="neg-prompt-input")
            generate_button = gr.Button("Generate", elem_id="generate-button", scale=0.13)

        with gr.Row():
            with gr.Column():
                with gr.Row():
                    scheduler_input = gr.Dropdown(choices=schedulers, label="Sampling method", value=schedulers[0])
                    seed_input = gr.Number(label="Seed", value=-1)

                with gr.Row():
                    width_input = gr.Slider(minimum=128, maximum=2048, step=128, label="Width", value=1024)
                    batch_size = gr.Slider(minimum=1, maximum=24, step=1, label="Batch size", value=1)

                with gr.Row():
                    height_input = gr.Slider(minimum=128, maximum=2048, step=128, label="Height", value=1024)
                    batch_count = gr.Slider(minimum=1, maximum=24, step=1, label="Batch Count", value=1)

                with gr.Row():
                    num_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Sampling steps", value=20)
                    cfg_scale_input = gr.Slider(minimum=1, maximum=20, step=1, label="CFG Scale", value=7)

                with gr.Accordion("Hires. fix", open=False):
                    use_hires = gr.Checkbox(label="Use Hires?", value=False, scale=0)
                    with gr.Row():
                        upscaler = gr.Dropdown(choices=schedulers, label="Upscaler", value=schedulers[0])
                        upscale_by = gr.Slider(minimum=1, maximum=8, step=1, label="Upscale by", value=2)
                    with gr.Row():
                        hires_steps = gr.Slider(minimum=1, maximum=50, step=1, label="Hires Steps", value=20)
                        denois_strength = gr.Slider(minimum=0, maximum=1, step=0.02, label="Denoising Strength", value=2)

            with gr.Column():
                # Gallery untuk output gambar
                output_gallery = gr.Gallery(label="Image Results")
                # Output teks JSON di bawah gallery
                output_text = gr.Textbox(label="Metadata", placeholder="Results are in Json format", lines=2)

        def update_images(prompt, neg_prompt, width, height, scheduler, num_steps, batch_size, batch_count, cfg_scale, seed, model, vae):
            # Update fungsi sesuai kebutuhan
            return generate_image_wrapper(prompt, neg_prompt, width, height, scheduler, num_steps, batch_size, batch_count, cfg_scale, seed, model, vae)

        generate_button.click(
            fn=update_images, 
            inputs=[prompt_input, neg_prompt_input, width_input, height_input, scheduler_input, num_steps_input, batch_size, batch_count, cfg_scale_input, seed_input, checkpoint_dropdown, vae_dropdown], 
            outputs=[output_gallery, output_text]
        )
        
    # Tab untuk Image-to-Image
    with gr.Tab("Image-to-Image"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input_i2i = gr.Textbox(label="Prompt", placeholder="Masukkan prompt teks", lines=2, elem_id="prompt-input")
                neg_prompt_input_i2i = gr.Textbox(label="Neg Prompt", placeholder="Masukkan negasi prompt", lines=2, elem_id="neg-prompt-input")
            generate_button = gr.Button("Generate", elem_id="generate-button", scale=0.13)
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Unggah Gambar")
                generate_button_i2i = gr.Button("Generate")
                with gr.Row():
                    scheduler_input = gr.Dropdown(choices=schedulers, label="Sampling method", value=schedulers[0])
                    seed_input = gr.Number(label="Seed", value=-1)
                
                with gr.Row():
                    steps = gr.Slider(minimum=1, maximum=100, step=1, label="Steps", value=20)
                    cfg_scale = gr.Slider(minimum=1, maximum=24, step=1, label="CFG Scale", value=7)
                    
                with gr.Row():
                    strength = gr.Slider(minimum=0, maximum=1, step=0.1, label="Strength", value=0.6)
            
            with gr.Column():
                output_image_i2i = gr.Image(label="Hasil Gambar")
        
        def process_image_func(image, prompt, neg_prompt, model):
            # Update fungsi sesuai kebutuhan
            return process_image(image, prompt, neg_prompt, model)
        
        generate_button_i2i.click(fn=process_image_func, inputs=[image_input, prompt_input_i2i, neg_prompt_input_i2i, checkpoint_dropdown, vae_dropdown], outputs=output_image_i2i)

    # Tab untuk ControlNet
    with gr.Tab("ControlNet"):
        with gr.Row():
            with gr.Column():
                controlnet_dropdown = gr.Dropdown(choices=controlnet_types, label="Pilih Tipe ControlNet")
                controlnet_image_input = gr.Image(label="Unggah Gambar untuk ControlNet")
                controlnet_button = gr.Button("Proses dengan ControlNet")
            
            with gr.Column():
                controlnet_output_image = gr.Image(label="Hasil ControlNet")
        controlnet_button.click(fn=controlnet_process_func, inputs=[controlnet_image_input, controlnet_dropdown, checkpoint_dropdown, vae_dropdown], outputs=controlnet_output_image)

    # Tab untuk Intpainting
    with gr.Tab ("Inpainting"):
        with gr.Row():
            with gr.Column():
                image = gr.ImageMask(sources=["upload"], layers=False, transforms=[], format="png", label="base image", show_label=True)
                btn = gr.Button("Inpaint!", elem_id="run_button")
                prompt = gr.Textbox(placeholder="Your prompt (what you want in place of what is erased)", show_label=False, elem_id="prompt")
                negative_prompt = gr.Textbox(label="negative_prompt", placeholder="Your negative prompt", info="what you don't want to see in the image")
                guidance_scale = gr.Number(value=7.5, minimum=1.0, maximum=20.0, step=0.1, label="guidance_scale")
                steps = gr.Number(value=20, minimum=10, maximum=30, step=1, label="steps")
                strength = gr.Number(value=0.99, minimum=0.01, maximum=1.0, step=0.01, label="strength")
                scheduler = gr.Dropdown(label="Schedulers", choices=schedulers, value="EulerDiscreteScheduler")
            with gr.Column():
                image_out = gr.Image(label="Output", elem_id="output-img")
                   
    btn.click(fn=intpaint_func, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out])
    


    
    # Tab untuk Describe
    with gr.Tab("Describe"):
        with gr.Row():
            with gr.Column():
                # Components
                image = gr.Image(type="pil", image_mode="RGBA", label="Input")
                submit_button = gr.Button(value="Submit", variant="primary", size="lg")
                model_repo = gr.Dropdown(modules.wdtagger.dropdown_list, value=modules.wdtagger.dropdown_list[0], label="Model")
                general_thresh = gr.Slider(0, 1, step=modules.wdtagger.args.score_slider_step, value=modules.wdtagger.args.score_general_threshold, label="General Tags Threshold", scale=3)
                general_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
                character_thresh = gr.Slider(0, 1, step=modules.wdtagger.args.score_slider_step, value=modules.wdtagger.args.score_character_threshold, label="Character Tags Threshold", scale=3)
                character_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
                clear_button = gr.ClearButton(components=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled], variant="secondary", size="lg")
                
            with gr.Column():
                sorted_general_strings = gr.Textbox(label="Output (string)")
                rating = gr.Label(label="Rating")
                character_res = gr.Label(label="Output (characters)")
                general_res = gr.Label(label="Output (tags)")
        
        clear_button.add([sorted_general_strings, rating, character_res, general_res])
        submit_button.click(modules.wdtagger.predictor.predict, inputs=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled], outputs=[sorted_general_strings, rating, character_res, general_res])

# Jalankan antarmuka
app.launch()