Spaces:
Running
on
Zero
Running
on
Zero
import time | |
import torch | |
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase | |
import gradio as gr | |
from gradio.components import Component | |
try: | |
import spaces | |
except ImportError: | |
class spaces: | |
def GPU(*args, **kwargs): | |
return lambda x: x | |
from output import UpsamplingOutput | |
from utils import IMAGE_SIZE_OPTIONS, RATING_OPTIONS, LENGTH_OPTIONS, IDENTITY_OPTIONS | |
ALL_MODELS = { | |
"dart-v2-llama-100m-sft": { | |
"repo": "p1atdev/dart-v2-llama-100m-sft", | |
"type": "sft", | |
}, | |
"dart-v2-mistral-100m-sft": { | |
"repo": "p1atdev/dart-v2-mistral-100m-sft", | |
"type": "sft", | |
}, | |
"dart-v2-mixtral-160m-sft": { | |
"repo": "p1atdev/dart-v2-mixtral-160m-sft", | |
"type": "sft", | |
}, | |
} | |
def prepare_models(model_name: str): | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_name, | |
torch_dtype=torch.bfloat16, | |
device_map="auto", | |
) | |
return { | |
"tokenizer": tokenizer, | |
"model": model, | |
} | |
def normalize_tags(tokenizer: PreTrainedTokenizerBase, tags: str): | |
"""Just remove unk tokens.""" | |
return ", ".join( | |
tokenizer.batch_decode( | |
[ | |
token | |
for token in tokenizer.encode_plus( | |
tags.strip(), | |
return_tensors="pt", | |
).input_ids[0] | |
if int(token) != tokenizer.unk_token_id | |
], | |
skip_special_tokens=True, | |
) | |
) | |
def compose_prompt( | |
copyright: str = "", | |
character: str = "", | |
general: str = "", | |
rating: str = "<|rating:sfw|>", | |
aspect_ratio: str = "<|aspect_ratio:tall|>", | |
length: str = "<|length:long|>", | |
identity: str = "<|identity:none|>", | |
): | |
prompt = ( | |
f"<|bos|>" | |
f"<copyright>{copyright.strip()}</copyright>" | |
f"<character>{character.strip()}</character>" | |
f"{rating}{aspect_ratio}{length}" | |
f"<general>{general.strip()}{identity}<|input_end|>" | |
) | |
return prompt | |
def generate_tags( | |
model, | |
tokenizer: PreTrainedTokenizerBase, | |
prompt: str, | |
): | |
print( # debug | |
tokenizer.tokenize( | |
prompt, | |
add_special_tokens=False, | |
) | |
) | |
input_ids = tokenizer.encode_plus(prompt, return_tensors="pt").input_ids | |
output = model.generate( | |
input_ids.to(model.device), | |
do_sample=True, | |
temperature=1, | |
top_p=0.9, | |
top_k=100, | |
num_beams=1, | |
num_return_sequences=1, | |
max_length=256, | |
) | |
# remove input tokens | |
pure_output_ids = output[0][len(input_ids[0]) :] | |
return ", ".join( | |
[ | |
token | |
for token in tokenizer.batch_decode( | |
pure_output_ids, skip_special_tokens=True | |
) | |
if token.strip() != "" | |
] | |
) | |
class V2UI: | |
model_name: str | None = None | |
model: AutoModelForCausalLM | |
tokenizer: PreTrainedTokenizerBase | |
input_components: list[Component] = [] | |
generate_btn: gr.Button | |
def on_generate( | |
self, | |
model_name: str, | |
copyright_tags: str, | |
character_tags: str, | |
general_tags: str, | |
rating_option: str, | |
# aspect_ratio_option: str, | |
length_option: str, | |
identity_option: str, | |
image_size: str, # this is from image generation config | |
*args, | |
) -> UpsamplingOutput: | |
if self.model_name is None or self.model_name != model_name: | |
models = prepare_models(ALL_MODELS[model_name]["repo"]) | |
self.model = models["model"] | |
self.tokenizer = models["tokenizer"] | |
self.model_name = model_name | |
# normalize tags | |
copyright_tags = normalize_tags(self.tokenizer, copyright_tags) | |
character_tags = normalize_tags(self.tokenizer, character_tags) | |
general_tags = normalize_tags(self.tokenizer, general_tags) | |
rating_tag = RATING_OPTIONS[rating_option] | |
aspect_ratio_tag = IMAGE_SIZE_OPTIONS[image_size] | |
length_tag = LENGTH_OPTIONS[length_option] | |
identity_tag = IDENTITY_OPTIONS[identity_option] | |
prompt = compose_prompt( | |
copyright=copyright_tags, | |
character=character_tags, | |
general=general_tags, | |
rating=rating_tag, | |
aspect_ratio=aspect_ratio_tag, | |
length=length_tag, | |
identity=identity_tag, | |
) | |
start = time.time() | |
upsampled_tags = generate_tags( | |
self.model, | |
self.tokenizer, | |
prompt, | |
) | |
elapsed_time = time.time() - start | |
return UpsamplingOutput( | |
upsampled_tags=upsampled_tags, | |
copyright_tags=copyright_tags, | |
character_tags=character_tags, | |
general_tags=general_tags, | |
rating_tag=rating_tag, | |
aspect_ratio_tag=aspect_ratio_tag, | |
length_tag=length_tag, | |
identity_tag=identity_tag, | |
elapsed_time=elapsed_time, | |
) | |
def ui(self): | |
input_copyright = gr.Textbox( | |
label="Copyright tags", | |
placeholder="vocaloid", | |
) | |
input_character = gr.Textbox( | |
label="Character tags", | |
placeholder="hatsune miku", | |
) | |
input_general = gr.TextArea( | |
label="General tags", | |
lines=4, | |
placeholder="1girl, ...", | |
value="1girl", | |
) | |
input_rating = gr.Radio( | |
label="Rating", | |
choices=list(RATING_OPTIONS.keys()), | |
value="general", | |
) | |
# input_aspect_ratio = gr.Radio( | |
# label="Aspect ratio", | |
# choices=["ultra_wide", "wide", "square", "tall", "ultra_tall"], | |
# value="tall", | |
# ) | |
input_length = gr.Radio( | |
label="Length", | |
choices=list(LENGTH_OPTIONS.keys()), | |
value="long", | |
) | |
input_identity = gr.Radio( | |
label="Identity", | |
choices=list(IDENTITY_OPTIONS.keys()), | |
value="lax", | |
) | |
model_name = gr.Dropdown( | |
label="Model", | |
choices=list(ALL_MODELS.keys()), | |
value=list(ALL_MODELS.keys())[0], | |
) | |
self.generate_btn = gr.Button(value="Generate", variant="primary") | |
self.input_components = [ | |
model_name, | |
input_copyright, | |
input_character, | |
input_general, | |
input_rating, | |
# input_aspect_ratio, | |
input_length, | |
input_identity, | |
] | |
def get_generate_btn(self) -> gr.Button: | |
return self.generate_btn | |
def get_inputs(self) -> list[Component]: | |
return self.input_components | |