Update app.py

app.py

@@ -6,6 +6,8 @@ from diffusers import StableDiffusionPipeline
 from transformers import CLIPTokenizer
 import os
 import zipfile
+import tempfile
+import shutil
 import gradio as gr
 
 # Define the device
@@ -104,9 +106,24 @@ def zip_model(model_path):
 
 # Gradio interface functions
 def start_fine_tuning(uploaded_files, prompts, num_epochs):
-    images = [Image.open(file).convert("RGB") for file in uploaded_files]
+    # Create a temporary directory for storing files
+    temp_dir = tempfile.mkdtemp()
+    print("Temporary directory:", temp_dir)
+
+    images = []
+    for file in uploaded_files:
+        # Store the uploaded file in the temp directory
+        image_path = os.path.join(temp_dir, file.name)
+        with open(image_path, 'wb') as f:
+            f.write(file.read())  # Save file content
+        images.append(Image.open(image_path).convert("RGB"))
+    
     model_save_path = "fine_tuned_model"
     fine_tune_model(images, prompts, model_save_path, num_epochs=int(num_epochs))
+    
+    # Clean up the temporary directory after fine-tuning
+    shutil.rmtree(temp_dir)
+
     return "Fine-tuning completed! Model is ready for download."
 
 def download_model():
@@ -156,4 +173,4 @@ with gr.Blocks() as demo:
 
         generate_button.click(generate_new_image, [prompt_input], generated_image)
 
-demo.launch()
+demo.launch()