Spaces:
Running
on
Zero
Running
on
Zero
import concurrent.futures | |
import random | |
import gradio as gr | |
import requests, os | |
import io, base64, json | |
import spaces | |
import torch | |
from PIL import Image | |
from openai import OpenAI | |
from .models import IMAGE_GENERATION_MODELS, VIDEO_GENERATION_MODELS, load_pipeline | |
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 | |
class ModelManager: | |
def __init__(self): | |
self.model_ig_list = IMAGE_GENERATION_MODELS | |
self.model_ie_list = IMAGE_EDITION_MODELS | |
self.model_vg_list = VIDEO_GENERATION_MODELS | |
self.loaded_models = {} | |
def load_model_pipe(self, model_name): | |
if not model_name in self.loaded_models: | |
pipe = load_pipeline(model_name) | |
self.loaded_models[model_name] = pipe | |
else: | |
pipe = self.loaded_models[model_name] | |
return pipe | |
def generate_image_ig(self, prompt, model_name): | |
pipe = self.load_model_pipe(model_name) | |
if 'Stable-cascade' not in model_name: | |
result = pipe(prompt=prompt).images[0] | |
else: | |
prior, decoder = pipe | |
prior.enable_model_cpu_offload() | |
prior_output = prior( | |
prompt=prompt, | |
height=512, | |
width=512, | |
negative_prompt='', | |
guidance_scale=4.0, | |
num_images_per_prompt=1, | |
num_inference_steps=20 | |
) | |
decoder.enable_model_cpu_offload() | |
result = decoder( | |
image_embeddings=prior_output.image_embeddings.to(torch.float16), | |
prompt=prompt, | |
negative_prompt='', | |
guidance_scale=0.0, | |
output_type="pil", | |
num_inference_steps=10 | |
).images[0] | |
return result | |
def generate_image_ig_api(self, prompt, model_name): | |
pipe = self.load_model_pipe(model_name) | |
result = pipe(prompt=prompt) | |
return result | |
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 == "": | |
from .matchmaker import matchmaker | |
not_run = [] #12,13,14,15,16,17,18,19,20,21,22, 25,26 #23,24, | |
model_ids = matchmaker(num_players=len(self.model_ig_list), not_run=not_run) | |
print(model_ids) | |
model_names = [self.model_ig_list[i] for i in model_ids] | |
print(model_names) | |
else: | |
model_names = [model_A, model_B, model_C, model_D] | |
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] | |
return results[0], results[1], results[2], results[3], \ | |
model_names[0], model_names[1], model_names[2], model_names[3] | |
def generate_image_ig_cache_anony(self, model_A, model_B, model_C, model_D): | |
if model_A == "" and model_B == "" and model_C == "" and model_D == "": | |
from .matchmaker import matchmaker | |
not_run = [20,21,22] | |
model_ids = matchmaker(num_players=len(self.model_ig_list), not_run=not_run) | |
print(model_ids) | |
model_names = [self.model_ig_list[i] for i in model_ids] | |
print(model_names) | |
else: | |
model_names = [model_A, model_B, model_C, model_D] | |
root_dir = SSH_CACHE_IMAGE | |
local_dir = "./cache_image" | |
if not os.path.exists(local_dir): | |
os.makedirs(local_dir) | |
prompt, results = get_ssh_random_image_prompt(root_dir, local_dir, model_names) | |
return results[0], results[1], results[2], results[3], \ | |
model_names[0], model_names[1], model_names[2], model_names[3], prompt | |
def generate_video_vg_parallel_anony(self, model_A, model_B, model_C, model_D): | |
if model_A == "" and model_B == "" and model_C == "" and model_D == "": | |
# model_names = random.sample([model for model in self.model_vg_list], 4) | |
from .matchmaker_video import matchmaker_video | |
model_ids = matchmaker_video(num_players=len(self.model_vg_list)) | |
print(model_ids) | |
model_names = [self.model_vg_list[i] for i in model_ids] | |
print(model_names) | |
else: | |
model_names = [model_A, model_B, model_C, model_D] | |
root_dir = SSH_CACHE_OPENSOURCE | |
for name in model_names: | |
if "Runway-Gen3" in name or "Runway-Gen2" in name or "Pika-v1.0" in name: | |
root_dir = SSH_CACHE_ADVANCE | |
elif "Pika-beta" in name: | |
root_dir = SSH_CACHE_PIKA | |
elif "Sora" in name and "OpenSora" not in name: | |
root_dir = SSH_CACHE_SORA | |
local_dir = "./cache_video" | |
if not os.path.exists(local_dir): | |
os.makedirs(local_dir) | |
prompt, results = get_ssh_random_video_prompt(root_dir, local_dir, model_names) | |
cache_dir = local_dir | |
return results[0], results[1], results[2], results[3], \ | |
model_names[0], model_names[1], model_names[2], model_names[3], prompt, cache_dir | |
def generate_image_ig_museum_parallel_anony(self, model_A, model_B, model_C, model_D): | |
if model_A == "" and model_B == "" and model_C == "" and model_D == "": | |
# model_names = random.sample([model for model in self.model_ig_list], 4) | |
from .matchmaker import matchmaker | |
model_ids = matchmaker(num_players=len(self.model_ig_list)) | |
print(model_ids) | |
model_names = [self.model_ig_list[i] for i in model_ids] | |
print(model_names) | |
else: | |
model_names = [model_A, model_B, model_C, model_D] | |
prompt = get_random_mscoco_prompt() | |
print(prompt) | |
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] | |
return results[0], results[1], results[2], results[3], \ | |
model_names[0], model_names[1], model_names[2], model_names[3], prompt | |
def generate_image_ig_parallel(self, prompt, model_A, model_B): | |
model_names = [model_A, model_B] | |
with concurrent.futures.ThreadPoolExecutor() as executor: | |
futures = [executor.submit(self.generate_image_ig, prompt, model) if model.startswith("imagenhub") | |
else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names] | |
results = [future.result() for future in futures] | |
return results[0], results[1] | |
def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name): | |
pipe = self.load_model_pipe(model_name) | |
result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct) | |
return result | |
def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B): | |
model_names = [model_A, model_B] | |
with concurrent.futures.ThreadPoolExecutor() as executor: | |
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] | |
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 == "": | |
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: | |
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] | |