diff --git a/skynet/cli.py b/skynet/cli.py index 9da80de..8800981 100755 --- a/skynet/cli.py +++ b/skynet/cli.py @@ -97,6 +97,45 @@ def inpaint(model, prompt, input, mask, output, strength, guidance, steps, seed) seed=seed ) + +@click.command() +@click.option('--model', '-m', default='microsoft/DialoGPT-small') +@click.option( + '--prompt', '-p', default='a red old tractor in a sunny wheat field') +@click.option('--output', '-o', default='output.txt') +@click.option('--temperature', '-t', default=1.0) +@click.option('--max-length', '-ml', default=256) +@click.option('--num-return-sequences', '-rs', defautl=1) +@click.option('--no-repeat-ngram', '-nr', default=2) +@click.option('--top-p', '-tp', default=0.95) +def txt2txt( + model, + prompt, + output, + temperature, + max_length, + num_return_sequences, + no_repeat_ngram, + top_p +): + from . import utils + config = load_skynet_toml() + hf_token = load_key(config, 'skynet.dgpu.hf_token') + hf_home = load_key(config, 'skynet.dgpu.hf_home') + set_hf_vars(hf_token, hf_home) + utils.txt2txt( + hf_token=hf_token, + model=model, + prompt=prompt, + output=output, + temperature=temperature, + max_length=max_length, + num_return_sequences=num_return_sequences, + no_repeat_ngram=no_repeat_ngram, + top_p=top_p + ) + + @click.command() @click.option('--input', '-i', default='input.png') @click.option('--output', '-o', default='output.png') diff --git a/skynet/dgpu/compute.py b/skynet/dgpu/compute.py index 049f00b..878cebd 100644 --- a/skynet/dgpu/compute.py +++ b/skynet/dgpu/compute.py @@ -8,7 +8,8 @@ import logging from hashlib import sha256 import zipfile from PIL import Image -from diffusers import DiffusionPipeline +from diffusers import DiffusionPipeline, AutoPipelineForInpainting +from transformers import AutoModelForCausalLM, AutoTokenizer import trio import torch @@ -63,6 +64,19 @@ def prepare_params_for_diffuse( ) +def prepare_params_for_transform( + params: dict, +): + return ( + params['prompt'], + int(params['num_return_sequences']), + int(params['no_repeat_ngram_size']), + float(params['top_p']), + float(params['temperature']), + int(params['max_length']), + ) + + class SkynetMM: def __init__(self, config: dict): @@ -100,7 +114,9 @@ class SkynetMM: self._model = pipeline_for( name, mode, cache_dir=self.cache_dir) - def get_model(self, name: str, mode: str) -> DiffusionPipeline: + def get_model( + self, name: str, mode: str + ) -> DiffusionPipeline | AutoPipelineForInpainting | AutoModelForCausalLM: if name not in MODELS: raise DGPUComputeError(f'Unknown model {model_name}') @@ -167,6 +183,26 @@ class SkynetMM: output_hash = sha256(output_binary).hexdigest() + case 'txt2txt': + arguments = prepare_params_for_transform(params) + prompt, num_return_sequences, no_repeat_ngram_size, top_p, temperature, max_length = arguments + tokenizer = AutoTokenizer.from_pretrained(model) + + self.get_model(params['model'], method) + tokenized_input = tokenizer.encode(prompt, return_tensors='pt') + tokenized_output = self._model.generate( + tokenized_input, + max_length=max_length, + num_return_sequences=num_return_sequences, + no_repeat_ngram_size=no_repeat_ngram_size, + top_p=top_p, + temperature=temperature, + ) + + output = tokenizer.decode(tokenized_output, skip_special_tokens=True) + + output_hash = sha256(output).hexdigest() + case _: raise DGPUComputeError('Unsupported compute method') diff --git a/skynet/utils.py b/skynet/utils.py index 0b43451..0993b78 100755 --- a/skynet/utils.py +++ b/skynet/utils.py @@ -22,6 +22,7 @@ from diffusers import ( EulerAncestralDiscreteScheduler ) from realesrgan import RealESRGANer +from transformers import AutoTokenizer, AutoModelForCausalLM from huggingface_hub import login import trio @@ -68,7 +69,7 @@ def pipeline_for( mode: str, mem_fraction: float = 1.0, cache_dir: str | None = None -) -> DiffusionPipeline: +) -> DiffusionPipeline | AutoPipelineForInpainting | AutoModelForCausalLM: assert torch.cuda.is_available() torch.cuda.empty_cache() @@ -109,6 +110,10 @@ def pipeline_for( if 'inpaint' in mode: pipe_class = AutoPipelineForInpainting + elif 'txt2txt' in mode: + # TODO: unsure about pipe + pipe_class = AutoModelForCausalLM + else: pipe_class = DiffusionPipeline @@ -233,6 +238,36 @@ def inpaint( image.save(output) +def txt2txt( + hf_token: str, + model: str = 'tiiuae/falcon-40b-instruct', + prompt: str = 'who is ur faja?', + output: str = 'output.txt', + num_return_sequences: int = 1, + no_repeat_ngram_size: int = 2, + top_p: float = 0.95, + temperature: float = 1.0, + max_length: int = 256, +): + login(token=hf_token) + tokenizer = AutoTokenizer.from_pretrained(model) + pipe = pipeline_for(model, 'txt2txt') + + prompt = prompt + tokenized_input = tokenizer.encode(prompt, return_tensors='pt') + tokenized_output = pipe.generate( + tokenized_input, + max_length=max_length, + num_return_sequences=num_return_sequences, + no_repeat_ngram_size=no_repeat_ngram_size, + top_p=top_p, + temperature=temperature + ) + response = tokenizer.decode(tokenized_output, skip_special_tokens=True) + with open(output, 'w', encoding='utf-8') as f: + f.write(response) + + def init_upscaler(model_path: str = 'weights/RealESRGAN_x4plus.pth'): return RealESRGANer( scale=4,