Update app.py

app.py

@@ -1,4 +1,46 @@
-from main import create_gradio_ui
+# app.py
 
-if __name__ == "__main__":
-    create_gradio_ui()
+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()