Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,323 Bytes
f7a6c5f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
from .imagen_museum import TASK_DICT, DOMAIN
from .imagen_museum import fetch_indexes
import random
ARENA_TO_IG_MUSEUM = {"LCM(v1.5/XL)":"LCM",
"PlayGroundV2.5": "PlayGroundV2_5"}
def draw2_from_imagen_museum(task, model_name1, model_name2):
task_name = TASK_DICT[task]
model_name1 = ARENA_TO_IG_MUSEUM[model_name1] if model_name1 in ARENA_TO_IG_MUSEUM else model_name1
model_name2 = ARENA_TO_IG_MUSEUM[model_name2] if model_name2 in ARENA_TO_IG_MUSEUM else model_name2
domain = DOMAIN
baselink = domain + task_name
matched_results = fetch_indexes(baselink)
r = random.Random()
uid, value = r.choice(list(matched_results.items()))
image_link_1 = baselink + "/" + model_name1 + "/" + uid
image_link_2 = baselink + "/" + model_name2 + "/" + uid
if task == "t2i": # Image Gen
prompt = value['prompt']
return [[image_link_1, image_link_2], [prompt]]
if task == "tie": # Image Edit
instruction = value['instruction']
input_caption = value['source_global_caption']
output_caption = value['target_global_caption']
source_image_link = baselink + "/" + "input" + "/" + uid
return [[source_image_link, image_link_1, image_link_2], [input_caption, output_caption, instruction]]
else:
raise ValueError("Task not supported")
def draw_from_imagen_museum(task, model_name):
task_name = TASK_DICT[task]
model_name = ARENA_TO_IG_MUSEUM[model_name] if model_name in ARENA_TO_IG_MUSEUM else model_name
domain = DOMAIN
baselink = domain + task_name
matched_results = fetch_indexes(baselink)
r = random.Random()
uid, value = r.choice(list(matched_results.items()))
model = model_name
image_link = baselink + "/" + model + "/" + uid
print(image_link)
if task == "t2i": # Image Gen
prompt = value['prompt']
return [image_link, prompt]
if task == "tie": # Image Edit
instruction = value['instruction']
input_caption = value['source_global_caption']
output_caption = value['target_global_caption']
source_image_link = baselink + "/" + "input" + "/" + uid
return [[source_image_link, image_link], [input_caption, output_caption, instruction]]
else:
raise ValueError("Task not supported")
|