feat: Fix selected space loading in app.py

The code changes fix the loading of the selected space in the `infer` function of `app.py`. Previously, the code was not properly handling the case when the selected space failed to load. The updated code now attempts to load the next available space from the list of spaces until a successful load is achieved. This ensures that the model is properly loaded and prevents any errors from occurring during inference.

This commit builds upon the recent user commits that updated the selected space and integrated the FLUX.1 [schnell] model with Gradio. It addresses the issue of failed space loading and improves the reliability of the model integration.

app.py

@@ -7,24 +7,34 @@ MAX_IMAGE_SIZE = 2048
 
 
 flux_1_schell_spaces = ["https://black-forest-labs-flux-1-schnell.hf.space", "ChristianHappy/FLUX.1-schnell", "innoai/FLUX.1-schnell", "tuan2308/FLUX.1-schnell", "FiditeNemini/FLUX.1-schnell"]
-flux_1_schnell_space = "https://black-forest-labs-flux-1-schnell.hf.space"
+# flux_1_schnell_space = "https://black-forest-labs-flux-1-schnell.hf.space"
 
 client = None
 job = None
 
-selected_space = gr.State(flux_1_schell_spaces[0]);
+selected_space_index = gr.State(0);
 
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+def infer(selected_space_index, prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     global job
     global client
-    if client is None:
+
+    # Asegúrate de que selected_space_index esté inicializado antes de este bloque de código
+    max_attempts = len(flux_1_schell_spaces)
+    attempts = 0
+
+    while client is None and attempts < max_attempts:
         try:
-            client = Client(flux_1_schnell_space)
-            print(f"Loaded custom model from {flux_1_schnell_space}")
+            selected_space = flux_1_schell_spaces[selected_space_index]
+            client = Client(selected_space)
+            print(f"Loaded custom model from {selected_space}")
         except ValueError as e:
-            print(f"Failed to load custom model: {e}")
+            print(f"Failed to load custom model from {selected_space}: {e}")
+            selected_space_index = (selected_space_index + 1) % len(flux_1_schell_spaces)
             client = None
-            raise gr.Error("Failed to load client for " + flux_1_schnell_space)
+            attempts += 1
+
+    if client is None:
+        raise gr.Error("Failed to load client after trying all spaces.")
         
     try:
         job = client.submit(
@@ -40,7 +50,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     except ValueError as e:
         raise gr.Error(e)
     
-    return result
+    return (selected_space_index, ) + result
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -130,8 +140,8 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs = [result, seed]
+        inputs = [selected_space_index, prompt, seed, randomize_seed, width, height, num_inference_steps],
+        outputs = [selected_space_index, result, seed]
     )
 
 demo.launch()