Update app.py

app.py

@@ -22,15 +22,16 @@ import torch
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 from diffusers import DiffusionPipeline
+import spaces
 
-
+@spaces.GPU()
 def save_state_to_file(state):
     filename = "state.pkl"
     with open(filename, "wb") as f:
         pickle.dump(state, f)
     return filename
 
-
+@spaces.GPU()
 def load_state_from_file(filename):
     with open(filename, "rb") as f:
         state = pickle.load(f)
@@ -47,10 +48,10 @@ blur_value = None
 allowed_res_max = 1.0
 
 # Device configuration
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+device = "cuda"
 print(f"Using device: {device}")
 
-
+@spaces.GPU()
 def weight_population(layer_type, resolution, depth, value):
     # Check if layer_type exists, if not, create it
     if layer_type not in weights:
@@ -67,6 +68,7 @@ def weight_population(layer_type, resolution, depth, value):
     # Add/Modify the value at the specified depth (which can be a string)
     weights[layer_type][resolution][depth] = value
 
+@spaces.GPU()
 def resize_image_with_aspect(image, res_range_min=128, res_range_max=1024):
     # Get the original width and height of the image
     width, height = image.size
@@ -89,6 +91,7 @@ def resize_image_with_aspect(image, res_range_min=128, res_range_max=1024):
     
     return resized_image
 
+@spaces.GPU()
 def reconstruct(input_img, caption):
     global weights
     weights = {}
@@ -187,6 +190,7 @@ def reconstruct(input_img, caption):
 
         return image_np, caption, 12, [caption, real_image_initial_latents.detach(), inversed_latents, weights]
 
+@spaces.GPU()
 class AttnReplaceProcessor(AttnProcessor2_0):
 
     def __init__(self, replace_all, layer_type, layer_count, blur_sigma=None):
@@ -294,6 +298,7 @@ class AttnReplaceProcessor(AttnProcessor2_0):
 
         return hidden_states
 
+@spaces.GPU()
 def replace_attention_processor(unet, clear=False, blur_sigma=None):
     attention_count = 0
 
@@ -314,7 +319,7 @@ def replace_attention_processor(unet, clear=False, blur_sigma=None):
             else:
                 module.processor = AttnReplaceProcessor(False, layer_type, attention_count, blur_sigma=blur_sigma)
 
-
+@spaces.GPU()
 def apply_prompt(meta_data, new_prompt):
 
     caption, real_image_initial_latents, inversed_latents, _ = meta_data
@@ -362,7 +367,7 @@ def apply_prompt(meta_data, new_prompt):
 
     return image_np
 
-
+@spaces.GPU()
 def on_image_change(filepath):
     # Extract the filename without extension
     filename = os.path.splitext(os.path.basename(filepath))[0]
@@ -395,7 +400,7 @@ def on_image_change(filepath):
 
     return filepath, img, meta_data_raw, num_inference_steps, scale_value, scale_value
 
-
+@spaces.GPU()
 def update_value(value, layer_type, resolution, depth):
     global weights
     weights[layer_type][resolution][depth] = value
@@ -422,6 +427,7 @@ def adjust_ends(values, adjustment):
 
 max_scale_value = 16
 
+@spaces.GPU()
 def update_scale(scale):
     global weights