Views
No views yet
1
2from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
3import torch
4
5def main():
6 #////////////////////////////////////////////
7 seed = 42
8 model = "Norod78/sd2-dreambooth-ClaymationXmas"
9 #////////////////////////////////////////////
10
11 torch.manual_seed(seed)
12 generator = torch.Generator()
13 generator.manual_seed(seed)
14
15 scheduler = DPMSolverMultistepScheduler(
16 beta_start=0.00085,
17 beta_end=0.012,
18 beta_schedule="scaled_linear",
19 num_train_timesteps=1000,
20 trained_betas=None,
21 predict_epsilon=True,
22 thresholding=False,
23 algorithm_type="dpmsolver++",
24 solver_type="midpoint",
25 lower_order_final=True,
26)
27
28 device = "cuda" if torch.cuda.is_available() else "cpu"
29 dtype = torch.float16 if device == "cuda" else torch.float32
30 pipe = StableDiffusionPipeline.from_pretrained(model, scheduler=scheduler,torch_dtype=dtype, generator=generator,use_auth_token=True).to(device)
31
32 #////////////////////////////////////////////
33 num_inference_steps = 20
34 width=512
35 height=512
36 samples=4
37 #////////////////////////////////////////////
38
39 prompt = "Willy Wonka, ClaymationXmas"
40 result = pipe([prompt] * samples, num_inference_steps=num_inference_steps, height=height, width=width)
41 images = result["images"]
42 for i, image in enumerate(images):
43 prompt_to_print = str(i) + "-" + prompt
44 output_file = prompt_to_print.replace(" ", "_") + "-" + str(width) + "x" +str(height)+ "_" + str(num_inference_steps) + "steps" + "_seed" + str(seed) + ".jpg"
45 image.save(output_file)
46 print("Saved: " + str(output_file))
47
48if __name__ == '__main__':
49 main()