| from diffusers import StableDiffusionPipeline |
| import torch |
|
|
|
|
| class Model: |
| def __init__(self, checkpoint="checkpoints/lora30k", weight_name="pytorch_lora_weights.safetensors", device="cuda"): |
| self.checkpoint = checkpoint |
| state_dict, network_alphas = StableDiffusionPipeline.lora_state_dict( |
| |
| checkpoint, |
| weight_name=weight_name |
| ) |
| self.pipe = StableDiffusionPipeline.from_pretrained( |
| "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16).to(device) |
| self.pipe.load_lora_into_unet(state_dict, network_alphas, self.pipe.unet, adapter_name='test_lora') |
| self.pipe.load_lora_into_text_encoder(state_dict, network_alphas, self.pipe.text_encoder, adapter_name='test_lora') |
| self.pipe.set_adapters(["test_lora"], adapter_weights=[1.0]) |
|
|
|
|
| def generate(self, prompt, negprompt='', steps=50, savedir=None, seed=1): |
| lora_scale = 1.0 |
| image = self.pipe(prompt, |
| negative_prompt=negprompt, |
| num_inference_steps=steps, |
| cross_attention_kwargs={"scale": lora_scale}, |
| generator=torch.manual_seed(seed)).images[0] |
| if savedir is None: |
| image.save(f"{self.checkpoint}/{'_'.join(prompt.replace('.', ' ').split(' '))}.png") |
| else: |
| image.save(f"{savedir}/{'_'.join(prompt.replace('.', ' ').split(' '))}.png") |
| return image |
|
|
|
|
| if __name__ == "__main__": |
|
|
| model = Model() |
|
|
| prompt = 'A happy 55 year old male with blond hair and a goatee. Visible teeth.' |
| negprompt = '' |
|
|
| image = model.generate(prompt, negprompt=negprompt, steps=50, seed=42) |
|
|
