How to Run and Host Flux.1 Image Generation (on your own cloud)

$ pip install "runhouse[aws]" 
$ aws configure 

import runhouse as rh

# First, we define a class that will hold the model and allow us to send prompts to it.
# We'll later wrap this with `rh.module`. This is a Runhouse class that allows you to
# run code in your class on a remote machine.
#
# Learn more at run.house/docs/tutorials/api-modules.
class FluxPipeline:
    def __init__(
        self,
        model_id: str = "black-forest-labs/FLUX.1-schnell",  # Schenll is smaller and faster while dev is more powerful but slower
    ):
        super().__init__()
        self.model_id = model_id
        self.pipeline = None

    def _load_pipeline(self):
        import torch
        from diffusers import FluxPipeline

        if not self.pipeline:
            self.pipeline = FluxPipeline.from_pretrained(
                self.model_id, torch_dtype=torch.bfloat16, use_safetensors=True
            )
            self.pipeline.enable_sequential_cpu_offload()  # Optimizes memory usage to allow the model to fit and inference on an A10 which has 24GB of memory

    def generate(self, input_prompt: str, **parameters):
        import torch

        torch.cuda.empty_cache()

        if not self.pipeline:
            self._load_pipeline()

        image = self.pipeline(
            input_prompt,
            guidance_scale=0.0,
            num_inference_steps=4,
            max_sequence_length=256,
            generator=torch.Generator("cpu").manual_seed(0),
        ).images[0]

        return image

# Now, we define the main function that will run locally when we run this script, and set up our Runhouse module on a remote cluster. First, we create a cluster with the desired instance type and provider. Our `instance_type` here is defined as `g5.8xlarge`, which is an AWS instance type. We can alternatively specify an accelerator type and count, such as `A10G:1`, and any instance type with those specifications will be used.
if __name__ == "__main__":

    cluster = rh.cluster(
        name="rh-g5",
        instance_type="g5.8xlarge",
        provider="aws",
    ).up_if_not()

    # Next, we define the environment for our module. This includes the required dependencies that need to be installed on the remote machine, as well as any secrets that need to be synced up from local to remote.
    env = rh.env(
        name="flux_inference",
        reqs=[
            "diffusers",
            "torch",
            "transformers[sentencepiece]",
            "accelerate",
        ],
    )

    # Finally, we define our module and run it on the remote cluster. We construct it normally and then call `to` to run it on the remote cluster. Alternatively, we could first check for an existing instance on the cluster by calling `cluster.get(name="flux")`. This would return the remote model after an initial run.
    RemoteFlux = rh.module(FluxPipeline).to(cluster, env=env, name="FluxPipeline")
    remote_flux = RemoteFlux(
        name="flux"
    )  # This has now been set up as a service on the remote cluster and can be used for inference.

    # Calling our remote function
    # We can call the `generate` method on the model class instance if it were running locally. This will run the function on the remote cluster and return the response to our local machine automatically. Further calls will also run on the remote machine, and maintain state that was updated between calls
    prompt = "A woman runs through a large, grassy field towards a house."
    response = remote_flux.generate(prompt)
    response.save("flux-schnell.png")
    response.show()