Update app.py

app.py

@@ -88,7 +88,6 @@ def get_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
 
     return normals
 
-
 @spaces.GPU
 def generate(
     prompt: str,
@@ -203,22 +202,18 @@ def launch_app(gen_fn: Callable):
     demo.queue().launch(share=False)
 
 
-def main():
-    # Initialize the models and dataset
-    diffuser, clip_model, dataset, device = init("config")
-    generate_sample = partial(
-        generate,
-        dataset=dataset,
-        device=device,
-        diffuser=diffuser,
-        clip_model=clip_model,
-    )
-
-    launch_app(generate_sample)
-
-
 # ------------------------------------------------------------------------------------- #
 
+diffuser, clip_model, dataset, device = init("config")
+diffuser.to(device)
+clip_model.to(device)
+
+generate_sample = partial(
+    generate,
+    dataset=dataset,
+    device=device,
+    diffuser=diffuser,
+    clip_model=clip_model,
+)
 
-if __name__ == "__main__":
-    main()
+launch_app(generate_sample)