Load model weights without Hugging Face or S3
data/
data/ scheduler/ scheduler_config.json text_encoder/ config.json diffusion_pytorch_model.bin tokenizer/ merges.txt tokenizer_config.json vocab.json unet/ config.json diffusion_pytorch_model.bin vae/ config.json diffusion_pytorch_model.bin model_index.json
model.py
class Model: def __init__(self, **kwargs): self._data_dir = kwargs["data_dir"] def load(self): self.model = StableDiffusionPipeline.from_pretrained( str(self._data_dir), revision="fp16", torch_dtype=torch.float16, ).to("cuda")
config.yaml
secrets: aws_access_key_id: null aws_secret_access_key: null aws_region: null # e.g., us-east-1 aws_bucket: null
import boto3 def __init__(self, **kwargs): self._config = kwargs.get("config") secrets = kwargs.get("secrets") self.s3_client = boto3.client( "s3", aws_access_key_id=secrets["aws_access_key_id"], aws_secret_access_key=secrets["aws_secret_access_key"], region_name=secrets["aws_region"], ) self.s3_bucket = secrets["aws_bucket"]
truss push
Was this page helpful?