Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import os
 import shutil
-from pathlib import Path
 from main import fine_tune_model
 from diffusers import StableDiffusionPipeline, DDIMScheduler
 import torch
@@ -9,31 +8,41 @@ import torch
 MODEL_NAME = "runwayml/stable-diffusion-v1-5"
 OUTPUT_DIR = "/home/user/app/stable_diffusion_weights/custom_model"
 
-def fine_tune(instance_prompt, image1, image2):
+def fine_tune(instance_prompt, image1, image2=None):
     instance_data_dir = "/home/user/app/instance_images"
-    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."
+    
+    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 generate_images(prompt, num_samples, height, width, num_inference_steps, guidance_scale):
-    pipe = StableDiffusionPipeline.from_pretrained(OUTPUT_DIR, safety_checker=None, torch_dtype=torch.float16).to("cuda")
-    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-    g_cuda = torch.Generator(device='cuda').manual_seed(1337)
-
-    with torch.autocast("cuda"), 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, generator=g_cuda
-        ).images
-
-    return images
+    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.float16).to("cuda")
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        g_cuda = torch.Generator(device='cuda').manual_seed(1337)
+        
+        with torch.autocast("cuda"), 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, generator=g_cuda
+            ).images
+        
+        return images
+    except Exception as e:
+        return str(e)
 
 def gradio_app():
     with gr.Blocks() as demo:
@@ -64,4 +73,4 @@ def gradio_app():
     demo.launch()
 
 if __name__ == "__main__":
-    gradio_app()
+    gradio_app()