Update app.py

app.py

@@ -142,8 +142,7 @@ def generate(
 # ------------------------------------------------------------------------------------- #
 
 
-@spaces.GPU
-def main(gen_fn: Callable):
+def launch_app(gen_fn: Callable):
     theme = gr.themes.Default(primary_hue="blue", secondary_hue="gray")
 
     with gr.Blocks(theme=theme) as demo:
@@ -203,10 +202,8 @@ def main(gen_fn: Callable):
     demo.queue().launch(share=False)
 
 
-# ------------------------------------------------------------------------------------- #
-
-
-if __name__ == "__main__":
+@spaces.GPU
+def main():
     # Initialize the models and dataset
     diffuser, clip_model, dataset, device = init("config")
     generate_sample = partial(
@@ -217,4 +214,11 @@ if __name__ == "__main__":
         clip_model=clip_model,
     )
 
-    main(generate_sample)
+    launch_app(generate_sample)
+
+
+# ------------------------------------------------------------------------------------- #
+
+
+if __name__ == "__main__":
+    main()