runhouse_1.py

runouse_1.py

import runhouse as rh
from diffusers import StableDiffusionPipeline

MODEL_ID = "stabilityai/stable-diffusion-2-base"

# My function
def sd_generate(prompt, **inference_kwargs):
    model = StableDiffusionPipeline.from_pretrained(MODEL_ID)
    model = model.to("cuda")
    return model(prompt, **inference_kwargs).images


if __name__ == "__main__":
    # The compute
    gpu = rh.cluster(name="rh-a10x-a",
                     instance_type="A10G:1",
                     provider="aws").up_if_not()

    # The environment, its own process on the cluster
    sd_env = rh.env(reqs=["torch", "transformers", "diffusers", "accelerate"])

    # Deploying my function to my env on the cluster
    remote_sd_generate = rh.function(sd_generate).to(gpu, env=sd_env)

    # Calling my function normally, but it's running remotely
    imgs = remote_sd_generate("A hot dog made out of matcha.")
    imgs[0].show()