Update app.py

app.py

@@ -7,7 +7,7 @@ from huggingface_hub import login, HfApi, HfFolder
 from PIL import Image
 import os
 from datetime import datetime
-
+import shutil
 
 # Get your Hugging Face API token
 folder = HfFolder()
@@ -33,27 +33,31 @@ pipe2.enable_xformers_memory_efficient_attention()
 
 def save_image_to_hf_space(image_np, image_name):
     # Name of your Hugging Face repo
-    repo_name = "FFusion/FFusionXL-SDXL-DEMO"
+    repo_name = "FFusion/FF2"
 
     # Convert the numpy array to an image
     image = Image.fromarray(image_np.astype('uint8'))
 
     # Append a timestamp to the filename
     timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
-    image_name_with_timestamp = f"{image_name}-{timestamp}"
+    image_name_with_timestamp = f"{image_name}-{timestamp}.png"
 
     # Save the image locally
-    image_path = f"./{image_name_with_timestamp}.png"
-    image.save(image_path)
+    local_image_path = f"./{image_name_with_timestamp}"
+    image.save(local_image_path)
 
     # Upload the image to your Hugging Face repo
     api.upload_file(
         token=token,
-        path_or_fileobj=image_path,
+        path_or_fileobj=local_image_path,
         repo_id=repo_name,
-        filename=image_name_with_timestamp,
+        path_in_repo=image_name_with_timestamp  # The path where the image will be stored in the repository
     )
 
+    # Save the image to the persistent storage
+    persistent_storage_path = f"/data/{image_name_with_timestamp}"
+    shutil.copy(local_image_path, persistent_storage_path)
+
     
 def genie (prompt, negative_prompt, scale, steps, seed):
     torch.cuda.empty_cache()
@@ -65,8 +69,8 @@ def genie (prompt, negative_prompt, scale, steps, seed):
     refined_image_np = np.array(refined_images[0])
 
     # Save the generated images to Hugging Face Spaces
-    # save_image_to_hf_space(int_image_np, "int_image")
-    # save_image_to_hf_space(refined_image_np, "refined_image")
+    save_image_to_hf_space(int_image_np, "int_image")
+    save_image_to_hf_space(refined_image_np, "refined_image")
 
     return int_image_np, refined_image_np