handle input validation outside GPU

to show the error message in the UI, otherwise generic error is shown: 'ZeroGPU worker error'

app.py

@@ -47,10 +47,8 @@ def load_model(model_class_name, model_filename, repo_id: str = "rizavelioglu/tr
     model.load_state_dict(state_dict, strict=True)
     return model.eval()
 
-@spaces.GPU(duration=10)
-@torch.no_grad()
-@timer_func
-def generate_multi_image(input_image, garment_types, seed=42, guidance_scale=2.0, num_inference_steps=50, is_upscale=False):
+def validate_garment_selection(garment_types):
+    """Validate garment type selection and return selected types and label indices."""
     label_map = {"Upper-Body": 0, "Lower-Body": 1, "Dress": 2}
     valid_single = ["Upper-Body", "Lower-Body", "Dress"]
     valid_tuple = ["Upper-Body", "Lower-Body"]
@@ -63,7 +61,19 @@ def generate_multi_image(input_image, garment_types, seed=42, guidance_scale=2.0
         selected, label_indices = valid_tuple, [label_map[t] for t in valid_tuple]
     else:
         raise gr.Error("Invalid selection. Choose one garment type or Upper-Body and Lower-Body together.")
+    
+    return selected, label_indices
+
+def generate_multi_image_wrapper(input_image, garment_types, seed=42, guidance_scale=2.0, num_inference_steps=50, is_upscale=False):
+    """Wrapper function that validates input before calling the GPU function."""
+    # Validate selection before entering GPU context
+    selected, label_indices = validate_garment_selection(garment_types)
+    return generate_multi_image(input_image, selected, label_indices, seed, guidance_scale, num_inference_steps, is_upscale)
 
+@spaces.GPU(duration=10)
+@torch.no_grad()
+@timer_func
+def generate_multi_image(input_image, selected, label_indices, seed=42, guidance_scale=2.0, num_inference_steps=50, is_upscale=False):
     batch_size = len(selected)
     scheduler.set_timesteps(num_inference_steps)
     generator = torch.Generator(device=device).manual_seed(seed)
@@ -263,10 +273,10 @@ def create_multi_tab():
                 submit_btn = gr.Button("Generate")
             with gr.Column():
                 output_image = gr.Image(type="pil", label="Generated Garment", height=384, width=384)
-        gr.Examples(examples=examples, inputs=[input_image, garment_type, seed, guidance_scale, inference_steps, upscale], outputs=output_image, fn=generate_multi_image, cache_examples=False, examples_per_page=2)
+        gr.Examples(examples=examples, inputs=[input_image, garment_type, seed, guidance_scale, inference_steps, upscale], outputs=output_image, fn=generate_multi_image_wrapper, cache_examples=False, examples_per_page=2)
         gr.Markdown(article)
         submit_btn.click(
-            fn=generate_multi_image,
+            fn=generate_multi_image_wrapper,
             inputs=[input_image, garment_type, seed, guidance_scale, inference_steps, upscale],
             outputs=output_image
         )