Spaces:
Running
on
Zero
Running
on
Zero
| 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") |