Update app.py

app.py

@@ -6,8 +6,6 @@ from diffusers import StableDiffusionPipeline
 from transformers import CLIPTokenizer
 import os
 import zipfile
-import tempfile
-import shutil
 import gradio as gr
 
 # Define the device
@@ -104,26 +102,18 @@ def zip_model(model_path):
                 zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), model_path))
     return zip_path
 
+# Function to save uploaded files
+def save_uploaded_file(uploaded_file, save_path):
+    # Open the file in binary write mode
+    with open(save_path, 'wb') as f:
+        f.write(uploaded_file.data)  # Use .data for the file content
+    return f"File saved at {save_path}"
+
 # Gradio interface functions
 def start_fine_tuning(uploaded_files, prompts, num_epochs):
-    # 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"))
-    
+    images = [Image.open(file).convert("RGB") for file in uploaded_files]
     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():
@@ -173,4 +163,4 @@ with gr.Blocks() as demo:
 
         generate_button.click(generate_new_image, [prompt_input], generated_image)
 
-demo.launch()
+demo.launch()