Inference: Stable Diffusion and FLAN-T5

gpu = rh.cluster(name='rh-a10x')

def sd_generate(prompt, num_images=1, steps=100, guidance_scale=7.5, model_id='stabilityai/stable-diffusion-2-base'):
    # imports must be defined inside the function for notebook environments
    from diffusers import StableDiffusionPipeline
    import torch

    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision='fp16').to('cuda')
    return pipe([prompt] * num_images, num_inference_steps=steps, guidance_scale=guidance_scale).images

generate_gpu = rh.function(fn=sd_generate).to(
    gpu,
    env=['./', 'torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117', 'diffusers', 'transformers']
)

# for outputting images later
!pip install ipyplot
import ipyplot

rh_prompt = 'A digital illustration of a woman running on the roof of a house.'
images = generate_gpu(rh_prompt, num_images=4, steps=50)

ipyplot.plot_images(images)

# save the function to be reusable later on
generate_gpu.save(name='sd_generate')

# By default, the GPU will terminate after 30 min of inactivity.
# To keep it up to reuse it for the rest of the tutorials
gpu.keep_warm()

# To stop the cluster after 10 min of inactivity
# gpu.keep_warm(autostop_mins=10)

# To terminate the cluster through runhouse. It can also be terminated directly through the cloud provider
# or by using the CLI commands `sky down gpu` or `sky down --all`
# gpu.teardown()

def causal_lm_generate(prompt, model_id='google/flan-t5-xl', **model_kwargs):
    import runhouse as rh
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

    (tokenizer, model) = rh.get_pinned_object(model_id) or (None, None)

    # pin to memory if it is not in memory yet
    if model is None:
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to('cuda')
        rh.pin_to_memory(model_id, (tokenizer, model))
    inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
    outputs = model.generate(**inputs, **model_kwargs)
    return tokenizer.batch_decode(outputs, skip_special_tokens=True)

flan_t5_generate = rh.function(fn=causal_lm_generate).to(gpu, env=['./'])

my_prompt = "A detailed oil painting of"
sequences = flan_t5_generate(my_prompt, max_new_tokens=100, min_length=20, temperature=2.0, repetition_penalty=3.0,
                              use_cache=False, do_sample=True, num_beams=3, num_return_sequences=4)

full_seqs = [my_prompt + " " + seq for seq in sequences]
for seq in full_seqs:
  print(seq)

# We can directly access the function by the name we saved it by, even on a different environment or run
saved_sd_generate = rh.function('sd_generate_pinned')
paintings = saved_sd_generate(full_seqs, num_images=1, steps=50)

ipyplot.plot_images(paintings)