Update ssh

model/matchmaker.py

@@ -6,9 +6,19 @@ import io, os
 import sys
 sys.path.append('../')
 from serve.constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_SKILL
+trueskill_env = TrueSkill()
 
+ssh_matchmaker_client = None
+sftp_matchmaker_client = None
 
-trueskill_env = TrueSkill()
+def create_ssh_matchmaker_client(server, port, user, password):
+    global ssh_matchmaker_client, sftp_matchmaker_client
+    ssh_matchmaker_client = paramiko.SSHClient()
+    ssh_matchmaker_client.load_system_host_keys()
+    ssh_matchmaker_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh_matchmaker_client.connect(server, port, user, password)
+
+    sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
 
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
@@ -25,29 +35,20 @@ def serialize_rating(rating):
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 
-def create_ssh_client(server, port, user, password):
-    ssh = paramiko.SSHClient()
-    ssh.load_system_host_keys()
-    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    ssh.connect(server, port, user, password)
-    return ssh
-
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    global sftp_matchmaker_client
     data = {
         'ratings': [serialize_rating(r) for r in ratings],
         'comparison_counts': comparison_counts.tolist(),
         'total_comparisons': total_comparisons
     }  
     json_data = json.dumps(data)
-    sftp = ssh.open_sftp()
-    with sftp.open(SSH_SKILL, 'w') as f:
+    with sftp_matchmaker_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 
 def load_json_via_sftp():
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-    sftp = ssh.open_sftp()
-    with sftp.open(SSH_SKILL, 'r') as f:
+    global sftp_matchmaker_client
+    with sftp_matchmaker_client.open(SSH_SKILL, 'r') as f:
         data = json.load(f)
     ratings = [deserialize_rating(r) for r in data['ratings']]
     comparison_counts = np.array(data['comparison_counts'])

model/model_manager.py

@@ -59,12 +59,6 @@ class ModelManager:
                        else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names]
             results = [future.result() for future in futures]
 
-        # with concurrent.futures.ThreadPoolExecutor() as executor:
-        #     futures = [executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names]
-        #     results = [future.result() for future in futures]
-
-        # results = [self.generate_image_ig(prompt, model) for model in model_names]
-
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3]
 
@@ -81,27 +75,12 @@ class ModelManager:
 
         prompt = get_random_mscoco_prompt()
         print(prompt)
-        # with concurrent.futures.ThreadPoolExecutor() as executor:
-        #     model_1 = model_names[0].split('_')[1]
-        #     model_2 = model_names[1].split('_')[1]
-        #     model_3 = model_names[2].split('_')[1]
-        #     model_4 = model_names[3].split('_')[1]
-
-        #     result_list = draw2_from_imagen_museum("t2i", model_1, model_2, model_3, model_4)
-        #     image_links = result_list[0]
-        #     prompt_list = result_list[1]
-        #     print(prompt_list[0])
 
         with concurrent.futures.ThreadPoolExecutor() as executor:
             futures = [executor.submit(self.generate_image_ig, prompt, model) if model.startswith("huggingface")
                        else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names]
             results = [future.result() for future in futures]
-        # with concurrent.futures.ThreadPoolExecutor() as executor:
-        #     futures = [executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names]
-        #     results = [future.result() for future in futures]
 
-        # results = [self.generate_image_ig_api(prompt, model) for model in model_names]
-        # results = [future.result() for future in futures]
         return results[0], results[1], results[2], results[3], \
             model_names[0], model_names[1], model_names[2], model_names[3], prompt
 

model/models/huggingface_models.py

@@ -6,13 +6,20 @@ import torch
 
 
 def load_huggingface_model(model_name, model_type):
+    # if model_name == "SD-turbo":
+    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
+    # elif model_name == "SDXL-turbo":
+    #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
+    # else:
+    #     raise NotImplementedError
+    # pipe = pipe.to("cuda")
     if model_name == "SD-turbo":
-        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
+        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")
     elif model_name == "SDXL-turbo":
-        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
+        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
     else:
         raise NotImplementedError
-    pipe = pipe.to("cuda")
+    pipe = pipe.to("cpu")
     return pipe
 
 

serve/Ksort.py

@@ -5,7 +5,7 @@ from .constants import KSORT_IMAGE_DIR
 from .constants import COLOR1, COLOR2, COLOR3, COLOR4
 from .vote_utils import save_any_image
 from .utils import disable_btn, enable_btn, invisible_btn
-from .upload import create_remote_directory, upload_image, upload_informance, upload_ssh_all
+from .upload import create_remote_directory, upload_ssh_all
 import json
 
 def reset_level(Top_btn):

serve/gradio_web.py

@@ -33,8 +33,11 @@ from .Ksort import (
     reset_vote_text,
     text_response_rank_igm,
 )
-from serve.upload import get_random_mscoco_prompt
+from .upload import get_random_mscoco_prompt, create_ssh_client
+from .update_skill import create_ssh_skill_client
+from model.matchmaker import create_ssh_matchmaker_client
 from functools import partial
+from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD
 
 def build_side_by_side_ui_anony(models):
     notice_markdown = """
@@ -50,6 +53,9 @@ def build_side_by_side_ui_anony(models):
 """
 
     model_list = models.model_ig_list
+    create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    create_ssh_skill_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
 
     state0 = gr.State()
     state1 = gr.State()

serve/update_skill.py

@@ -9,6 +9,17 @@ trueskill_env = TrueSkill()
 sys.path.append('../')
 from model.models import IMAGE_GENERATION_MODELS
 
+ssh_skill_client = None
+sftp_skill_client = None
+
+def create_ssh_skill_client(server, port, user, password):
+    global ssh_skill_client, sftp_skill_client
+    ssh_skill_client = paramiko.SSHClient()
+    ssh_skill_client.load_system_host_keys()
+    ssh_skill_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh_skill_client.connect(server, port, user, password)
+
+    sftp_skill_client = ssh_skill_client.open_sftp()
 
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
@@ -25,29 +36,21 @@ def serialize_rating(rating):
 def deserialize_rating(rating_dict):
     return trueskill_env.Rating(mu=rating_dict['mu'], sigma=rating_dict['sigma'])
 
-def create_ssh_client(server, port, user, password):
-    ssh = paramiko.SSHClient()
-    ssh.load_system_host_keys()
-    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    ssh.connect(server, port, user, password)
-    return ssh
 
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    global sftp_skill_client
     data = {
         'ratings': [serialize_rating(r) for r in ratings],
         'comparison_counts': comparison_counts.tolist(),
         'total_comparisons': total_comparisons
     }  
     json_data = json.dumps(data)
-    sftp = ssh.open_sftp()
-    with sftp.open(SSH_SKILL, 'w') as f:
+    with sftp_skill_client.open(SSH_SKILL, 'w') as f:
         f.write(json_data)
 
 def load_json_via_sftp():
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-    sftp = ssh.open_sftp()
-    with sftp.open(SSH_SKILL, 'r') as f:
+    global sftp_skill_client
+    with sftp_skill_client.open(SSH_SKILL, 'r') as f:
         data = json.load(f)
     ratings = [deserialize_rating(r) for r in data['ratings']]
     comparison_counts = np.array(data['comparison_counts'])

serve/upload.py

@@ -7,81 +7,51 @@ import json
 import random
 from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_MSCOCO
 
+ssh_client = None
+sftp_client = None
+
 def create_ssh_client(server, port, user, password):
-    ssh = paramiko.SSHClient()
-    ssh.load_system_host_keys()
-    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    ssh.connect(server, port, user, password)
-    return ssh
+    global ssh_client, sftp_client
+    ssh_client = paramiko.SSHClient()
+    ssh_client.load_system_host_keys()
+    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh_client.connect(server, port, user, password)
+
+    sftp_client = ssh_client.open_sftp()
 
 def get_image_from_url(image_url):
     response = requests.get(image_url)
     response.raise_for_status()  # success
     return Image.open(io.BytesIO(response.content))
 
-# def upload_image_via_sftp(ssh, image, remote_image_path):
-#     if isinstance(image, str):
-#         print("get url image")
-#         image = get_image_from_url(image)
-#     with ssh.open_sftp() as sftp:
-#         with io.BytesIO() as image_byte_stream:
-#             image.save(image_byte_stream, format='JPEG')
-#             image_byte_stream.seek(0)  
-#             sftp.putfo(image_byte_stream, remote_image_path)
-#         print(f"Successfully uploaded image to {remote_image_path}")
-        
-# def upload_json_via_sftp(ssh, data, remote_json_path):
-#     json_data = json.dumps(data, indent=4)
-#     with ssh.open_sftp() as sftp:
-#         with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
-#             sftp.putfo(json_byte_stream, remote_json_path)
-#         print(f"Successfully uploaded JSON data to {remote_json_path}")
-
-def upload_image(states, output_dir):
-    pass
-    # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-    # for i in range(len(states)):
-    #     output_file = os.path.join(output_dir, f"{i}.jpg")
-    #     upload_image_via_sftp(ssh, states[i].output, output_file)
-
-def upload_informance(data, data_path):
-    pass
-    # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-    # upload_json_via_sftp(ssh, data, data_path)
 def get_random_mscoco_prompt():
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-    sftp = ssh.open_sftp()
-    files = sftp.listdir(SSH_MSCOCO)
+    global sftp_client
+    files = sftp_client.listdir(SSH_MSCOCO)
     txt_files = [file for file in files if file.endswith('.txt')]
 
     selected_files = random.sample(txt_files, 1) # get one prompt
 
     for file in selected_files:
         remote_file_path = os.path.join(SSH_MSCOCO, file)
-        with sftp.file(remote_file_path, 'r') as f:
+        with sftp_client.file(remote_file_path, 'r') as f:
             content = f.read().decode('utf-8')
             print(f"Content of {file}:")
             print("\n")
-    sftp.close()
-    ssh.close()
-
     return content
 
 
 def create_remote_directory(remote_directory):
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
-
-    stdin, stdout, stderr = ssh.exec_command(f'mkdir -p {SSH_LOG}/{remote_directory}')
+    global ssh_client
+    stdin, stdout, stderr = ssh_client.exec_command(f'mkdir -p {SSH_LOG}/{remote_directory}')
     error = stderr.read().decode('utf-8')
     if error:
         print(f"Error: {error}")
     else:
         print(f"Directory {remote_directory} created successfully.")
-    # ssh.close()
     return f'{SSH_LOG}/{remote_directory}'
     
 def upload_ssh_all(states, output_dir, data, data_path):
-    ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    global sftp_client
     output_file_list = []
     image_list = []
     for i in range(len(states)):
@@ -89,17 +59,17 @@ def upload_ssh_all(states, output_dir, data, data_path):
         output_file_list.append(output_file)
         image_list.append(states[i].output)
 
-    with ssh.open_sftp() as sftp:
-        for i in range(len(output_file_list)):
-            if isinstance(image_list[i], str):
-                print("get url image")
-                image_list[i] = get_image_from_url(image_list[i])
-            with io.BytesIO() as image_byte_stream:
-                image_list[i].save(image_byte_stream, format='JPEG')
-                image_byte_stream.seek(0)  
-                sftp.putfo(image_byte_stream, output_file_list[i])
-            print(f"Successfully uploaded image to {output_file_list[i]}")
-        json_data = json.dumps(data, indent=4)
-        with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
-            sftp.putfo(json_byte_stream, data_path)
-        print(f"Successfully uploaded JSON data to {data_path}")
+# with sftp_client as sftp:
+    for i in range(len(output_file_list)):
+        if isinstance(image_list[i], str):
+            print("get url image")
+            image_list[i] = get_image_from_url(image_list[i])
+        with io.BytesIO() as image_byte_stream:
+            image_list[i].save(image_byte_stream, format='JPEG')
+            image_byte_stream.seek(0)  
+            sftp_client.putfo(image_byte_stream, output_file_list[i])
+        print(f"Successfully uploaded image to {output_file_list[i]}")
+    json_data = json.dumps(data, indent=4)
+    with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
+        sftp_client.putfo(json_byte_stream, data_path)
+    print(f"Successfully uploaded JSON data to {data_path}")