File size: 1,845 Bytes
eab86f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import torch
from diffusers import StableDiffusionImg2ImgPipeline, \
StableDiffusionPipeline
def check_cuda_device():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
return device
def get_the_model(device=None):
model_id = "stabilityai/stable-diffusion-2"
pipe = StableDiffusionPipeline.from_pretrained(model_id,
torch_dtype=torch.float16)
if device:
pipe.to(device)
else:
device = check_cuda_device()
pipe.to(device)
return pipe
def get_image_to_image_model(path=None, device=None):
model_id = "stabilityai/stable-diffusion-2"
if path:
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
path,
torch_dtype=torch.float16)
else:
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16)
if device:
if device == "cuda" or device == "cpu":
pipe.to(device)
else:
device = check_cuda_device()
pipe.to(device)
return pipe
def gen_initial_img(int_prompt):
model = get_the_model(None)
image = model(int_prompt, num_inference_steps=100).images[0]
return image
def generate_story(int_prompt, steps, iterations=133):
image_dic = {}
init_img = gen_initial_img(int_prompt)
img2img_model = get_image_to_image_model()
img = init_img
for idx, step in enumerate(steps):
print(f"step: {idx}")
print(step)
image = img2img_model(prompt=step, image=img, strength=0.75, guidance_scale=7.5,
num_inference_steps=iterations).images[0]
image_dic[idx] = {
"image": image,
"prompt": step
}
img = image
return init_img, image_dic
|