add txt2txt utils and cli, start on dgpu

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')

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')
 

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,