Dreambot / img_gen.py
tugot17's picture
Upload 3 files
eab86f7
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