app.py

# app.py

import os
import gradio as gr
from huggingface_hub import HfFolder

def launch_gradio_app(fine_tune_model, load_model, generate_images, push_to_huggingface, repo_name):
    with gr.Blocks() as demo:
        gr.Markdown("# Dreambooth App")
        
        with gr.Tab("Fine-tune Model"):
            with gr.Row():
                instance_images = gr.File(label="Instance Images", file_count="multiple")
                class_images = gr.File(label="Class Images", file_count="multiple")
            with gr.Row():
                instance_prompt = gr.Textbox(label="Instance Prompt")
                class_prompt = gr.Textbox(label="Class Prompt")
            with gr.Row():
                num_train_steps = gr.Number(label="Number of Training Steps", value=800)
                fine_tune_button = gr.Button("Fine-tune Model")
        
        with gr.Tab("Generate Images"):
            with gr.Row():
                prompt = gr.Textbox(label="Prompt")
                negative_prompt = gr.Textbox(label="Negative Prompt")
            with gr.Row():
                num_samples = gr.Number(label="Number of Samples", value=1)
                guidance_scale = gr.Number(label="Guidance Scale", value=7.5)
            with gr.Row():
                height = gr.Number(label="Height", value=512)
                width = gr.Number(label="Width", value=512)
                num_inference_steps = gr.Slider(label="Number of Inference Steps", value=50, minimum=1, maximum=100)
            generate_button = gr.Button("Generate Images")
            output_images = gr.Gallery()
        
        with gr.Tab("Push to Hugging Face"):
            push_button = gr.Button("Push Model to Hugging Face")
            huggingface_link = gr.Textbox(label="Hugging Face Model Link")
        
        fine_tune_button.click(fine_tune_model, inputs=[instance_images, class_images, instance_prompt, class_prompt, num_train_steps], outputs=huggingface_link)
        
        generate_button.click(generate_images, inputs=[prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale], outputs=output_images)
        
        push_button.click(push_to_huggingface, inputs=[HfFolder.path, repo_name], outputs=huggingface_link)
    
    demo.launch()