fixed gpu spaces?

app.py

@@ -98,7 +98,7 @@ def get_script_args_info(exclude_args=None):
         {
             "name": "--iterations",
             "type": "number",
-            "default": 4000,
+            "default": 6000,
             "help": "Number of optimization iterations",
         },
         {
@@ -159,7 +159,7 @@ def get_script_args_info(exclude_args=None):
         {
             "name": "--pruning_max_swaps",
             "type": "number",
-            "default": 20,
+            "default": 50,
             "precision": 0,
             "help": "Max number of swaps allowed after pruning",
         },
@@ -182,7 +182,7 @@ def get_script_args_info(exclude_args=None):
         {
             "name": "--learning_rate_warmup_fraction",
             "type": "slider",
-            "default": 0.2,
+            "default": 0.01,
             "min": 0.0,
             "max": 1.0,
             "step": 0.01,
@@ -214,7 +214,7 @@ def get_script_args_info(exclude_args=None):
         {
             "name": "--num_init_rounds",
             "type": "number",
-            "default": 8,
+            "default": 32,
             "precision": 0,
             "help": "Number of rounds to choose the starting height map from.",
         },
@@ -295,17 +295,31 @@ else:
 def run_autoforge_process(cmd, log_path):
     from joblib import parallel_backend
     cli_args = cmd[1:]
-    autoforge_main = importlib.import_module("autoforge.__main__")
+
+    # Import the library's high-level module which contains the main() entrypoint
+    # that uses the internal @spaces.GPU decorators (autoforge/auto_forge.py).
+    # Do NOT fall back to any other module; require the project module to be present.
+    try:
+        autoforge_main = importlib.import_module("autoforge.auto_forge")
+    except Exception as e:
+        # If the direct import fails, write the error to the log and return failure.
+        with open(log_path, "a", encoding="utf-8") as log_f:
+            log_f.write(f"\nERROR: Could not import autoforge.auto_forge: {e}\n")
+        return -1
 
     exit_code = 0
     with open(log_path, "w", buffering=1, encoding="utf-8") as log_f, \
             redirect_stdout(log_f), redirect_stderr(log_f), parallel_backend("threading", n_jobs=-1):
         try:
+            # Provide argv so the autoforge CLI parser sees the same arguments as before.
             sys.argv = ["autoforge"] + cli_args
+            # Call the library's main() function. It may call sys.exit(); capture that.
             autoforge_main.main()
         except SystemExit as e:
-            exit_code = e.code
+            # Preserve the exit code behavior from the subprocess-style invocation.
+            exit_code = e.code if isinstance(e.code, int) or e.code is None else 0
         except Exception as e:
+            # Log exceptions to the log file and return a failure code.
             log_f.write(f"\nERROR: {e}\n")
             exit_code = -1
 
@@ -672,7 +686,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     visible=False,
                 )
 
-    @spaces.GPU(duration=150)
     def execute_autoforge_script(
         current_filaments_df_state_val, input_image, *accordion_param_values
     ):
@@ -774,17 +787,45 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 self.exc = None
 
             def run(self):
-                try:
-                    self.returncode = run_autoforge_process(self.cmd, self.log_path)
-                except Exception as e:
-                    self.exc = e
-                    with open(self.log_path, "a", encoding="utf-8") as lf:
-                        lf.write(
-                            "\nERROR: {}. This usually means there was no GPU or the process took too long.\n".format(
-                                exc_text(e)
-                            )
-                        )
-                    self.returncode = -1
+                    try:
+                        # Import joblib's parallel_backend to match previous environment setup
+                        from joblib import parallel_backend
+
+                        # Import the autoforge high-level module which contains main().
+                        # No fallback allowed; require this module to exist in the package.
+                        try:
+                            autoforge_main = importlib.import_module("autoforge.auto_forge")
+                        except Exception as e:
+                            # Log and fail the worker if import fails.
+                            with open(self.log_path, "a", encoding="utf-8") as lf:
+                                lf.write(f"\nERROR: Could not import autoforge.auto_forge: {e}\n")
+                            self.exc = e
+                            self.returncode = -1
+                            return
+
+                        # Redirect stdout/stderr into the provided log file while running
+                        with open(self.log_path, "a", encoding="utf-8") as lf, \
+                                redirect_stdout(lf), redirect_stderr(lf), parallel_backend("threading", n_jobs=-1):
+                            try:
+                                # Ensure CLI-like argv for the autoforge parser
+                                sys.argv = ["autoforge"] + (self.cmd[1:] if len(self.cmd) > 1 else [])
+                                autoforge_main.main()
+                                # If main exits normally, treat as success (0)
+                                self.returncode = 0
+                            except SystemExit as e:
+                                # Preserve exit code behavior
+                                self.returncode = e.code if isinstance(e.code, int) or e.code is None else 0
+                            except Exception as e:
+                                # Write exception to log and mark failure
+                                lf.write(f"\nERROR: {e}\n")
+                                self.exc = e
+                                self.returncode = -1
+                    except Exception as e_outer:
+                        # If import or log file operations fail, persist exception and code
+                        self.exc = e_outer
+                        with open(self.log_path, "a", encoding="utf-8") as lf:
+                            lf.write("\nERROR: {}\n".format(exc_text(e_outer)))
+                        self.returncode = -1
 
         try:
             worker = Worker(command, log_file)