Fix bugs

app.py

@@ -55,10 +55,10 @@ SAMPLE_GRID = np.stack(SAMPLE_GRID, axis=-1).reshape(-1, 1, 2)
 SAMPLE_GRID = np.repeat(SAMPLE_GRID, LENGTH, axis=1) # [N, F, 2]
 
 @spaces.GPU
-def set_seed(seed):
+def set_seed(seed, device):
     np.random.seed(seed)
     torch.manual_seed(seed)
-    return torch.Generator(DEVICE).manual_seed(seed)
+    return torch.Generator(device).manual_seed(seed)
 
 @spaces.GPU
 def set_model(pretrained_model_path):
@@ -87,10 +87,11 @@ def upload_image(image):
     image = TRANSFORMS(image)
     return image
 
+@spaces.GPU
 def run(input_image, input_image_end, pretrained_model_path, seed,
         prompt, n_prompt, num_inference_steps, guidance_scale,
         *layer_args):
-    generator = set_seed(seed)
+    generator = set_seed(seed, DEVICE)
     global layer_tracking_points
     args_layer_tracking_points = [layer_tracking_points[i].value for i in range(LAYER_CAPACITY)]