Sorting modalities in generate_output() backend.py for consistent generations

src/backend.py

@@ -242,25 +242,27 @@ def generate_output(s2l1c_input, s2l2a_input, s1rtc_input, dem_input, num_infere
         gr.Warning("You need to remove some of the inputs that you would like to generate. If all modalities are known, there is nothing to generate.")
         return s2l1c_input, s2l2a_input, s1rtc_input, dem_input
 
-    images=[]
-    condition_modalities=[]
+    # Instead of collecting in UI order, create ordered dictionaries
+    input_images = {}
     if s2l1c_active:
-        images.append(s2l1c_input)
-        condition_modalities.append('s2_l1c')
+        input_images['s2_l1c'] = s2l1c_input
     if s2l2a_active:
-        images.append(s2l2a_input)
-        condition_modalities.append('s2_l2a')
+        input_images['s2_l2a'] = s2l2a_input
     if s1rtc_active:
-        images.append(s1rtc_input)
-        condition_modalities.append('s1_rtc')
+        input_images['s1_rtc'] = s1rtc_input
     if dem_active:
-        images.append(dem_input)
-        condition_modalities.append('dem')
-                                    
+        input_images['dem'] = dem_input
+    
+    condition_modalities = list(input_images.keys())
+    
+    # Sort modalities and collect images in the same order
+    sorted_modalities = sorted(condition_modalities, key=lambda x: ['dem', 's1_rtc', 's2_l1c', 's2_l2a'].index(x))
+    sorted_images = [input_images[mod] for mod in sorted_modalities]
+    
     imgs_out = custom_inference(
-        images=images,
+        images=sorted_images,
         generate_modalities=[el for el in ['s2_l1c', 's2_l2a', 's1_rtc', 'dem'] if el not in condition_modalities],
-        condition_modalities=condition_modalities,
+        condition_modalities=sorted_modalities,
         num_inference_steps=num_inference_steps_slider,
         seed=seed
     )
@@ -268,22 +270,11 @@ def generate_output(s2l1c_input, s2l2a_input, s1rtc_input, dem_input, num_infere
     output = []
 
     # Collect outputs
-    if s2l1c_active:
-        output.append(s2l1c_input)
-    else:
-        output.append(to_PIL(imgs_out['s2_l1c'][0]))
-    if s2l2a_active:
-        output.append(s2l2a_input)
-    else:
-        output.append(to_PIL(imgs_out['s2_l2a'][0]))
-    if s1rtc_active:
-        output.append(s1rtc_input)
-    else:
-        output.append(to_PIL(imgs_out['s1_rtc'][0]))
-    if dem_active:
-        output.append(dem_input)
-    else:
-        output.append(to_PIL(imgs_out['dem'][0]))
+    for modality in sorted_modalities:
+        if modality in input_images:
+            output.append(input_images[modality])
+        else:
+            output.append(to_PIL(imgs_out[modality][0]))
         
     return output