Image-generation / image_generation.py
Haris Ali
Remove scheduler selection bug
bcc5908
Raw
History Blame Contribute Delete
2.45 kB
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler, PNDMScheduler, DPMSolverMultistepScheduler, LMSDiscreteScheduler, HeunDiscreteScheduler
import torch
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
schedulers = {
"EulerDiscreteScheduler": EulerDiscreteScheduler,
"PNDMScheduler": PNDMScheduler,
"DPMSolverMultistepScheduler": DPMSolverMultistepScheduler,
"LMSDiscreteScheduler": LMSDiscreteScheduler,
"HeunDiscreteScheduler": HeunDiscreteScheduler,
}
# negative_prompt = "ugly, text, words, characters, tiling, poorly drawn hands, poorly drawn feet, \
# poorly drawn face, out of frame, extra limbs, disfigured, deformed, \
# body out of frame, bad anatomy, watermark, signature, cut off, low contrast, \
# underexposed, overexposed, bad art, beginner, amateur, distorted face."
def generate_image(prompt,
negative_prompt,
model_id,
scheduler_name,
inference_steps,
guidance_scale,
num_of_images):
pipe = StableDiffusionPipeline.from_pretrained(model_id)
if scheduler_name != "None":
scheduler = schedulers[scheduler_name].from_pretrained(model_id, subfolder="scheduler")
pipe.scheduler = scheduler
pipe = pipe.to(torch_device)
prompts = [prompt] * num_of_images
negative_prompts = [negative_prompt] * num_of_images
images = pipe(prompts, num_inference_steps=inference_steps, negative_prompt=negative_prompts, guidance_scale=float(guidance_scale)).images
return images
# def generate_image(prompt,
# negative_prompt,
# model_dropdown,
# scheduler_dropdown,
# inference_steps,
# guidance_scale):
# print(prompt)
# print(negative_prompt)
# print(model_dropdown)
# print(scheduler_dropdown)
# print(inference_steps)
# print(guidance_scale)
# from PIL import Image
# import numpy as np
# # Create four noise images
# noise_images = []
# for _ in range(4):
# # Generate random noise array
# width, height = 512, 512 # Adjust as desired
# noise = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
# noise_image = Image.fromarray(noise)
# noise_images.append(noise_image)
# return noise_images