Update app.py

app.py

@@ -1,51 +1,75 @@
 import gradio as gr
-from main import DreamboothApp
+import os
+import shutil
+from main import fine_tune_model
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+import torch
 
-app = DreamboothApp(model_path="stable_diffusion_weights")
+MODEL_NAME = "runwayml/stable-diffusion-v1-5"
+OUTPUT_DIR = "/home/user/app/stable_diffusion_weights/custom_model"
 
-def train(instance_images, instance_prompt, num_class_images, max_train_steps):
-    app.train(instance_data_dir="instance_data",
-              class_data_dir="class_data",
-              instance_prompt=instance_prompt,
-              class_prompt="photo of a person",
-              num_class_images=num_class_images,
-              max_train_steps=max_train_steps)
-    return "Training completed. Model is ready for inference."
+def fine_tune(instance_prompt, image1, image2=None):
+    instance_data_dir = "/home/user/app/instance_images"
+    
+    try:
+        if os.path.exists(instance_data_dir):
+            shutil.rmtree(instance_data_dir)
+        os.makedirs(instance_data_dir, exist_ok=True)
+        
+        image1.save(os.path.join(instance_data_dir, "instance_0.png"))
+        if image2 is not None:
+            image2.save(os.path.join(instance_data_dir, "instance_1.png"))
+        
+        fine_tune_model(instance_data_dir, instance_prompt, MODEL_NAME, OUTPUT_DIR)
+        return "Model fine-tuning complete."
+    except Exception as e:
+        return str(e)
 
-def inference(prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale, seed):
-    app.load_model()
-    images = app.inference(prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale, seed)
-    return images
+def generate_images(prompt, num_samples, height, width, num_inference_steps, guidance_scale):
+    try:
+        if not os.path.exists(OUTPUT_DIR):
+            return "The model path does not exist."
+        
+        pipe = StableDiffusionPipeline.from_pretrained(OUTPUT_DIR, safety_checker=None, torch_dtype=torch.float32).to("cpu")
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        
+        with torch.autocast("cpu"), torch.inference_mode():
+            images = pipe(
+                prompt, height=height, width=width, num_images_per_prompt=num_samples,
+                num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+            ).images
+        
+        return images
+    except Exception as e:
+        return str(e)
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Stable Diffusion Dreambooth")
-    with gr.Tab("Training"):
-        with gr.Row():
-            instance_images = gr.File(label="Upload Instance Images (5-10 images recommended)", file_count="multiple")
-            with gr.Column():
-                instance_prompt = gr.Textbox(label="Instance Prompt", placeholder="Enter the prompt for your instance images")
-                num_class_images = gr.Number(label="Number of Class Images", value=50)
-                max_train_steps = gr.Number(label="Maximum Training Steps", value=800)
-                train_button = gr.Button("Train Model")
-        train_output = gr.Textbox(label="Training Output")
-        train_button.click(train, inputs=[instance_images, instance_prompt, num_class_images, max_train_steps], outputs=train_output)
-
-    with gr.Tab("Inference"):
-        with gr.Row():
-            with gr.Column():
-                prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
-                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here (optional)")
-                with gr.Row():
+def gradio_app():
+    with gr.Blocks() as demo:
+        with gr.Tab("Fine-Tune Model"):
+            with gr.Row():
+                with gr.Column():
+                    instance_prompt = gr.Textbox(label="Instance Prompt")
+                    image1 = gr.Image(label="Upload Image 1", type="pil")
+                    image2 = gr.Image(label="Upload Image 2 (Optional)", type="pil")
+                    fine_tune_button = gr.Button("Fine-Tune Model")
+                    output_text = gr.Textbox(label="Output")
+                fine_tune_button.click(fine_tune, inputs=[instance_prompt, image1, image2], outputs=output_text)
+        
+        with gr.Tab("Generate Images"):
+            with gr.Row():
+                with gr.Column():
+                    prompt = gr.Textbox(label="Prompt")
                     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="Steps", value=50)
-                seed = gr.Number(label="Seed (optional)", value=0)
-                generate_button = gr.Button("Generate Images")
-            with gr.Column():
-                gallery = gr.Gallery(label="Generated Images")
-        generate_button.click(inference, inputs=[prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale, seed], outputs=gallery)
+                    num_inference_steps = gr.Slider(label="Steps", value=50, minimum=1, maximum=100)
+                    generate_button = gr.Button("Generate Images")
+                with gr.Column():
+                    gallery = gr.Gallery(label="Generated Images")
+                generate_button.click(generate_images, inputs=[prompt, num_samples, height, width, num_inference_steps, guidance_scale], outputs=gallery)
+        
+    demo.launch()
 
-demo.launch()
+if __name__ == "__main__":
+    gradio_app()