update image_cache

.gitignore

@@ -174,6 +174,6 @@ ksort-logs/
 cache_video/
 cache_image/
 
-/model/models/generate_image_cache.py
-/model/models/generate_video_cache.py
-/get_webvid_prompt.py
+model/models/generate_image_cache.py
+model/models/generate_video_cache.py
+get_webvid_prompt.py

get_webvid_prompt.py

@@ -1,32 +0,0 @@
-from datasets import load_dataset
-import pandas as pd
-import re
-# # Load the WebVid dataset
-# dataset = load_dataset('webvid', 'webvid-10m', split='train')
-# from datasets import load_dataset
-
-ds = load_dataset("TempoFunk/webvid-10M", cache_dir="/mnt/data/lizhikai/webvid/")
-v = ds['validation']['name']
-# 定义字符串长度的合理范围
-MIN_LENGTH = 30
-MAX_LENGTH = 300
-pattern = re.compile(r'^[a-zA-Z\s]+$')
-
-# 过滤掉空字符串和特别长特别短的字符串
-v = [s for s in v if len(s) >= MIN_LENGTH and len(s) <= MAX_LENGTH and pattern.match(s)]
-
-# 指定保存文件的路径
-file_path = 'webvid_prompt.txt'
-
-# 打开文件，以写入模式
-with open(file_path, 'w', encoding='utf-8') as file:
-    # 遍历列表中的每个字符串并写入文件
-    for item in v:
-        if '\n' in item:
-            continue
-        else:
-            file.write(item + '\n')
-
-print("字符串列表已成功保存到文件中。")
-
-

model/model_manager.py

@@ -8,7 +8,6 @@ import torch
 from PIL import Image
 from openai import OpenAI
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, load_pipeline
-from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum
 from serve.upload import get_random_mscoco_prompt, get_random_video_prompt, get_ssh_random_video_prompt, get_ssh_random_image_prompt
 from serve.constants import SSH_CACHE_OPENSOURCE, SSH_CACHE_ADVANCE, SSH_CACHE_PIKA, SSH_CACHE_SORA, SSH_CACHE_IMAGE
 
@@ -62,14 +61,6 @@ class ModelManager:
         
         return result
 
-    def generate_image_ig_museum(self, model_name):
-        model_name = model_name.split('_')[1]
-        result_list = draw_from_imagen_museum("t2i", model_name)
-        image_link = result_list[0]
-        prompt = result_list[1]
-
-        return image_link, prompt
-
 
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B, model_C, model_D):
         if model_A == "" and model_B == "" and model_C == "" and model_D == "":
@@ -174,15 +165,6 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
 
-    def generate_image_ig_museum_parallel(self, model_A, model_B):
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            model_1 = model_A.split('_')[1]
-            model_2 = model_B.split('_')[1]
-            result_list = draw2_from_imagen_museum("t2i", model_1, model_2)
-            image_links = result_list[0]
-            prompt_list = result_list[1]
-        return image_links[0], image_links[1], prompt_list[0]
-
 
     @spaces.GPU(duration=200)
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
@@ -190,14 +172,6 @@ class ModelManager:
         result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
         return result
 
-    def generate_image_ie_museum(self, model_name):
-        model_name = model_name.split('_')[1]
-        result_list = draw_from_imagen_museum("tie", model_name)
-        image_links = result_list[0]
-        prompt_list = result_list[1]
-        # image_links = [src, model]
-        # prompt_list = [source_caption, target_caption, instruction]
-        return image_links[0], image_links[1], prompt_list[0], prompt_list[1], prompt_list[2]
 
     def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         model_names = [model_A, model_B]
@@ -208,17 +182,6 @@ class ModelManager:
             results = [future.result() for future in futures]
         return results[0], results[1]
 
-    def generate_image_ie_museum_parallel(self, model_A, model_B):
-        model_names = [model_A, model_B]
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            model_1 = model_names[0].split('_')[1]
-            model_2 = model_names[1].split('_')[1]
-            result_list = draw2_from_imagen_museum("tie", model_1, model_2)
-            image_links = result_list[0]
-            prompt_list = result_list[1]
-            # image_links = [src, model_A, model_B]
-            # prompt_list = [source_caption, target_caption, instruction]
-        return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2]
 
     def generate_image_ie_parallel_anony(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         if model_A == "" and model_B == "":
@@ -229,21 +192,3 @@ class ModelManager:
             futures = [executor.submit(self.generate_image_ie, textbox_source, textbox_target, textbox_instruct, source_image, model) for model in model_names]
             results = [future.result() for future in futures]
         return results[0], results[1], model_names[0], model_names[1]
-
-    def generate_image_ie_museum_parallel_anony(self, model_A, model_B):
-        if model_A == "" and model_B == "":
-            model_names = random.sample([model for model in self.model_ie_list], 2)
-        else:
-            model_names = [model_A, model_B]
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            model_1 = model_names[0].split('_')[1]
-            model_2 = model_names[1].split('_')[1]
-            result_list = draw2_from_imagen_museum("tie", model_1, model_2)
-            image_links = result_list[0]
-            prompt_list = result_list[1]
-            # image_links = [src, model_A, model_B]
-            # prompt_list = [source_caption, target_caption, instruction]
-        return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2], model_names[0], model_names[1]
-
-
-        raise NotImplementedError

model/models/generate_image_cache.py

@@ -1,99 +0,0 @@
-from huggingface_models import load_huggingface_model
-from replicate_api_models import load_replicate_model
-from openai_api_models import load_openai_model
-from other_api_models import load_other_model
-import concurrent.futures 
-import os
-import io, time
-import requests
-import json
-from PIL import Image
-
-
-IMAGE_GENERATION_MODELS = [ 
-                            # 'replicate_SDXL_text2image',
-                            # 'replicate_SD-v3.0_text2image',
-                            # 'replicate_SD-v2.1_text2image',
-                            # 'replicate_SD-v1.5_text2image',
-                            # 'replicate_SDXL-Lightning_text2image',
-                            # 'replicate_Kandinsky-v2.0_text2image',
-                            # 'replicate_Kandinsky-v2.2_text2image',
-                            # 'replicate_Proteus-v0.2_text2image',
-                            # 'replicate_Playground-v2.0_text2image',
-                            # 'replicate_Playground-v2.5_text2image',
-                            # 'replicate_Dreamshaper-xl-turbo_text2image',
-                            # 'replicate_SDXL-Deepcache_text2image',
-                            # 'replicate_Openjourney-v4_text2image',
-                            # 'replicate_LCM-v1.5_text2image',
-                            # 'replicate_Realvisxl-v3.0_text2image',
-                            # 'replicate_Realvisxl-v2.0_text2image',
-                            # 'replicate_Pixart-Sigma_text2image',
-                            # 'replicate_SSD-1b_text2image',
-                            # 'replicate_Open-Dalle-v1.1_text2image',
-                            # 'replicate_Deepfloyd-IF_text2image',
-                            # 'huggingface_SD-turbo_text2image',
-                            # 'huggingface_SDXL-turbo_text2image',
-                            # 'huggingface_Stable-cascade_text2image',
-                            # 'openai_Dalle-2_text2image',
-                            # 'openai_Dalle-3_text2image',
-                            'other_Midjourney-v6.0_text2image',
-                            'other_Midjourney-v5.0_text2image',
-                            # "replicate_FLUX.1-schnell_text2image",
-                            # "replicate_FLUX.1-pro_text2image",
-                            # "replicate_FLUX.1-dev_text2image",
-                            ]
-
-Prompts = [
-    # 'An aerial view of someone walking through a forest alone in the style of Romanticism.',
-    # 'With dark tones and backlit resolution, this oil painting depicts a thunderstorm over a cityscape.',
-    # 'The rendering depicts a futuristic train station with volumetric lighting in an Art Nouveau style.',
-    # 'An Impressionist illustration depicts a river winding through a meadow.', # featuring a thick black outline
-    # 'Photo of a black and white picture of a person facing the sunset from a bench.',
-    # 'The skyline of a city is painted in bright, high-resolution colors.',
-    # 'A sketch shows two robots talking to each other, featuring a surreal look and narrow aspect ratio.',
-    # 'An abstract Dadaist collage in neon tones and 4K resolutions of a post-apocalyptic world.',
-    # 'With abstract elements and a rococo style, the painting depicts a garden in high resolution.',
-    # 'A picture of a senior man walking in the rain and looking directly at the camera from a medium distance.',
-]
-
-def load_pipeline(model_name):
-    model_source, model_name, model_type = model_name.split("_")
-    if model_source == "replicate":
-        pipe = load_replicate_model(model_name, model_type)
-    elif model_source == "huggingface":
-        pipe = load_huggingface_model(model_name, model_type)
-    elif model_source == "openai":
-        pipe = load_openai_model(model_name, model_type)
-    elif model_source == "other":
-        pipe = load_other_model(model_name, model_type)
-    else:
-        raise ValueError(f"Model source {model_source} not supported")
-    return pipe
-
-def generate_image_ig_api(prompt, model_name):
-        pipe = load_pipeline(model_name)
-        result = pipe(prompt=prompt)
-        return result
-
-save_names = []
-for name in IMAGE_GENERATION_MODELS:
-    model_source, model_name, model_type = name.split("_")
-    save_names.append(model_name)
-
-for i, prompt in enumerate(Prompts):
-    print("save the {} prompt".format(i+1))
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = [executor.submit(generate_image_ig_api, prompt, model) for model in IMAGE_GENERATION_MODELS]
-        results = [future.result() for future in futures]
-
-    root_dir = '/rscratch/zhendong/lizhikai/ksort/ksort_image_cache/'
-    save_dir = os.path.join(root_dir, f'output-{i+4}')
-    if not os.path.exists(save_dir):
-        os.makedirs(save_dir, exist_ok=True)
-    with open(os.path.join(save_dir, "prompt.txt"), 'w', encoding='utf-8') as file:
-        file.write(prompt)
-
-    for j, result in enumerate(results):
-        result = result.resize((512, 512))
-        file_path = os.path.join(save_dir, f'{save_names[j]}.jpg')
-        result.save(file_path, format="JPEG")

model/models/generate_video_cache.py

@@ -1,62 +0,0 @@
-
-
-file_path = '/home/lizhikai/webvid_prompt100.txt'
-str_list = []
-with open(file_path, 'r', encoding='utf-8') as file:
-    for line in file:
-        str_list.append(line.strip())
-        if len(str_list) == 100:
-            break
-
-def generate_image_ig_api(prompt, model_name):
-    model_source, model_name, model_type = model_name.split("_")
-    pipe = load_replicate_model(model_name, model_type)
-    result = pipe(prompt=prompt)
-    return result
-model_names = ['replicate_Zeroscope-v2-xl_text2video',
-                    #    'replicate_Damo-Text-to-Video_text2video',
-                        'replicate_Animate-Diff_text2video',
-                        'replicate_OpenSora_text2video',
-                        'replicate_LaVie_text2video',
-                        'replicate_VideoCrafter2_text2video',
-                        'replicate_Stable-Video-Diffusion_text2video',
-                        ]
-save_names = []
-for name in model_names:
-    model_source, model_name, model_type = name.split("_")
-    save_names.append(model_name)
-
-for i, prompt in enumerate(str_list):
-    print("save the {} prompt".format(i+1))
-    # if i+1 < 97:
-    #     continue
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = [executor.submit(generate_image_ig_api, prompt, model) for model in model_names]
-        results = [future.result() for future in futures]
-
-    root_dir = '/mnt/data/lizhikai/ksort_video_cache/'
-    save_dir = os.path.join(root_dir, f'cache_{i+1}')
-    if not os.path.exists(save_dir):
-        os.makedirs(save_dir, exist_ok=True)
-    with open(os.path.join(save_dir, "prompt.txt"), 'w', encoding='utf-8') as file:
-        file.write(prompt)
-
-    # 下载视频并保存
-    repeat_num = 5
-    for j, url in enumerate(results):
-        while 1:
-            time.sleep(1)
-            response = requests.get(url, stream=True)
-            if response.status_code == 200:
-                file_path = os.path.join(save_dir, f'{save_names[j]}.mp4')
-                with open(file_path, 'wb') as file:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        file.write(chunk)
-                print(f"视频 {j} 已保存到 {file_path}")
-                break
-            else:
-                repeat_num = repeat_num - 1
-                if repeat_num == 0:
-                    print(f"视频 {j} 保存失败")
-                    # raise ValueError("Video request failed.")
-                continue