diff --git a/img_styler/image_prompt/stable_diffusion.py b/img_styler/image_prompt/stable_diffusion.py new file mode 100644 index 0000000..d8939d1 --- /dev/null +++ b/img_styler/image_prompt/stable_diffusion.py @@ -0,0 +1,27 @@ +import torch +from torch import autocast +from PIL import Image +from io import BytesIO +from diffusers import StableDiffusionImg2ImgPipeline + + +def generate_image_with_prompt(input_img_path: str, prompt_txt: str = "Face portrait", + output_path: str=None, save: bool=True): + torch.cuda.empty_cache() + device = "cuda" + # License: https://huggingface.co/spaces/CompVis/stable-diffusion-license + model_path = "./models/stable_diffusion_v1_4" + pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_path, revision="fp16", + torch_dtype=torch.float16,) + pipe = pipe.to(device) + + # Open image + image_input = Image.open(input_img_path).convert("RGB") + init_image = image_input.resize((512, 512)) + + with autocast(device): + images = pipe(prompt=prompt_txt, init_image=init_image, strength=0.5, guidance_scale=7.5)["sample"] + + file_name = output_path + '/result.jpg' + images[0].save(file_name) + return file_name