Update app.py

app.py

@@ -88,18 +88,22 @@ def get_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
 
     return normals
 
+
 @spaces.GPU
 def generate(
     prompt: str,
     seed: int,
     guidance_weight: float,
     sample_label: str,
-    # ----------------------- ß#
+    # ----------------------- #
     dataset: MultimodalDataset,
     device: torch.device,
     diffuser: Diffuser,
     clip_model: clip.model.CLIP,
 ) -> Dict[str, Any]:
+    diffuser.to(device)
+    clip_model.to(device)
+
     # Set arguments
     set_random_seed(seed)
     diffuser.gen_seeds = np.array([seed])
@@ -206,9 +210,6 @@ def launch_app(gen_fn: Callable):
 # ------------------------------------------------------------------------------------- #
 
 diffuser, clip_model, dataset, device = init("config")
-diffuser.to("cuda")
-clip_model.to("cuda")
-
 generate_sample = partial(
     generate,
     dataset=dataset,
@@ -216,5 +217,4 @@ generate_sample = partial(
     diffuser=diffuser,
     clip_model=clip_model,
 )
-
 launch_app(generate_sample)