Update app.py

app.py

@@ -75,6 +75,9 @@ import os
 
 import torch
 import paramiko
+import socket
+import threading # NEW IMPORT
+import queue     # NEW IMPORT
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -91,6 +94,8 @@ FTP_HOST = os.getenv("FTP_HOST")
 FTP_USER = os.getenv("FTP_USER")
 FTP_PASS = os.getenv("FTP_PASS")
 FTP_DIR = os.getenv("FTP_DIR")
+FTP_HOST_FALLBACK = os.getenv("FTP_HOST_FALLBACK")
+FTP_DIR_FALLBACK = os.getenv("FTP_DIR_FALLBACK")
 
 def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         # adjust the batch_size of prompt_embeds according to guidance_scale
@@ -129,7 +134,88 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
       #  pipeline.scheduler._step_index = pipeline.num_timesteps * 0.9
     return callback_kwargs
 
+
+# --- WORKER FUNCTION FOR THREADING ---
+# This function contains the logic to connect to a single host.
+# It will be executed by each of our threads.
+def connect_worker(host, result_queue):
+    """Tries to connect to a single host and puts the successful transport object into the queue."""
+    transport = None
+    try:
+        transport = paramiko.Transport((host, 22))
+        # We still use the 5-second timeout for the handshake
+        transport.start_client(timeout=5)
+        transport.auth_password(username=FTP_USER, password=FTP_PASS)
+        
+        # If we reach here, the connection was successful.
+        # Put the result in the queue for the main thread to use.
+        print(f"✅ Connection to {host} succeeded first.")
+        result_queue.put(transport)
+
+    except (paramiko.SSHException, socket.timeout, EOFError) as e:
+        # This is an expected failure, just print a note.
+        print(f"ℹ️ Connection to {host} failed or was too slow: {e}")
+        if transport:
+            transport.close()
+    except Exception as e:
+        # Handle any other unexpected errors.
+        print(f"❌ Unexpected error connecting to {host}: {e}")
+        if transport:
+            transport.close()
+            
 def upload_to_ftp(filename):
+    """
+    Attempts to connect to two FTP hosts simultaneously and uses the first one that responds.
+    It now uses a corresponding directory for the primary and fallback hosts.
+    """
+    hosts = [FTP_HOST]
+    if FTP_HOST_FALLBACK:
+        hosts.append(FTP_HOST_FALLBACK)
+    
+    result_queue = queue.Queue()
+    threads = []
+    
+    print(f"--> Racing connections to {hosts} for uploading {filename}...")
+
+    for host in hosts:
+        thread = threading.Thread(target=connect_worker, args=(host, result_queue))
+        thread.daemon = True
+        thread.start()
+        threads.append(thread)
+
+    try:
+        winning_transport = result_queue.get(timeout=7)
+        
+        # --- THIS IS THE NEW LOGIC ---
+        # 1. Determine which host won the race.
+        winning_host = winning_transport.getpeername()[0]
+        
+        # 2. Select the correct destination directory based on the winning host.
+        #    If the fallback directory isn't specified, it safely defaults to the primary directory.
+        if winning_host == FTP_HOST:
+            destination_directory = FTP_DIR
+        else:
+            destination_directory = FTP_DIR_FALLBACK if FTP_DIR_FALLBACK else FTP_DIR
+
+        print(f"--> Proceeding with upload to {winning_host} in directory {destination_directory}...")
+
+        # 3. Construct the full destination path using the selected directory.
+        sftp = paramiko.SFTPClient.from_transport(winning_transport)
+        destination_path = os.path.join(destination_directory, os.path.basename(filename))
+        sftp.put(filename, destination_path)
+        
+        print(f"✅ Successfully uploaded {filename}.")
+        
+        sftp.close()
+        winning_transport.close()
+
+    except queue.Empty:
+        print("❌ Critical Error: Neither FTP host responded in time.")
+    except Exception as e:
+        print(f"❌ An unexpected error occurred during SFTP operation: {e}")
+
+
+def upload_to_ftp_old(filename):
     try:
         transport = paramiko.Transport((FTP_HOST, 22))
         destination_path=FTP_DIR+filename