Views
No views yet
[!IMPORTANT]
Make sure to upgrade diffusers to >= 0.28.0:pip install -U diffusers --upgradeIn addition make sure to installtransformers,safetensors,sentencepiece,gcandaccelerate:pip install transformers accelerate safetensors sentencepiece gc
1import torch
2from diffusers import StableDiffusionXLPipeline
3import gc
4from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection
5# from onediffx import compile_pipe, save_pipe, load_pipe
6from PIL import Image
7
8def encode_prompt(prompts, tokenizers, text_encoders):
9 embeddings_list = []
10 for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
11 cond_input = tokenizer(
12 prompt,
13 max_length=tokenizer.model_max_length,
14 padding='max_length',
15 truncation=True,
16 return_tensors='pt',
17 )
18 prompt_embeds = text_encoder(cond_input.input_ids.to('cuda'), output_hidden_states=True)
19 pooled_prompt_embeds = prompt_embeds[0]
20 embeddings_list.append(prompt_embeds.hidden_states[-2])
21 prompt_embeds = torch.concat(embeddings_list, dim=-1)
22 negative_prompt_embeds = torch.zeros_like(prompt_embeds)
23 negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
24 bs_embed, seq_len, _ = prompt_embeds.shape
25 prompt_embeds = prompt_embeds.repeat(1, 1, 1).view(bs_embed * 1, seq_len, -1)
26 seq_len = negative_prompt_embeds.shape[1]
27 negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1).view(1 * 1, seq_len, -1)
28 pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
29 negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
30 return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
31
32def get_user_input():
33 prompt = input("Enter prompt: ") or '3/4 shot, candid photograph of a beautiful 30 year old redhead woman with messy dark hair, peacefully sleeping in her bed, night, dark, light from window, dark shadows, masterpiece, uhd, moody'
34 negative_prompt = input("Enter negative prompt: ") or ""
35 cfg_scale = float(input("Enter CFG scale (default 7.5): ") or 7.5)
36 steps = int(input("Enter number of steps (default 50): ") or 50)
37 width = int(input("Enter image width (default 512): ") or 512)
38 height = int(input("Enter image height (default 512): ") or 512)
39 seed = int(input("Enter seed (default 42): ") or 42)
40 return prompt, negative_prompt, cfg_scale, steps, width, height, seed
41
42pipe = StableDiffusionXLPipeline.from_pretrained('ABDALLALSWAITI/DAVINCI-DIFF', use_safetensors=True, torch_dtype=torch.float16, tokenizer=None, text_encoder=None, tokenizer_2=None, text_encoder_2=None).to('cuda')
43# pipe = compile_pipe(pipe)
44# load_pipe(pipe, dir="cached_pipe")
45
46
47
48while True:
49 queue = []
50 prompt, negative_prompt, cfg_scale, steps, width, height, seed = get_user_input()
51 queue.append({
52 'prompt': prompt,
53 'negative_prompt': negative_prompt,
54 'cfg_scale': cfg_scale,
55 'steps': steps,
56 'width': width,
57 'height': height,
58 'seed': seed,
59 })
60
61 tokenizer = CLIPTokenizer.from_pretrained('ABDALLALSWAITI/DAVINCI-DIFF', subfolder='tokenizer')
62 text_encoder = CLIPTextModel.from_pretrained('ABDALLALSWAITI/DAVINCI-DIFF', subfolder='text_encoder', use_safetensors=True, torch_dtype=torch.float16).to('cuda')
63 tokenizer_2 = CLIPTokenizer.from_pretrained('ABDALLALSWAITI/DAVINCI-DIFF', subfolder='tokenizer_2')
64 text_encoder_2 = CLIPTextModelWithProjection.from_pretrained('ABDALLALSWAITI/DAVINCI-DIFF', subfolder='text_encoder_2', use_safetensors=True, torch_dtype=torch.float16).to('cuda')
65
66 with torch.no_grad():
67 for generation in queue:
68 generation['embeddings'] = encode_prompt(
69 [generation['prompt'], generation['prompt']],
70 [tokenizer, tokenizer_2],
71 [text_encoder, text_encoder_2],
72 )
73 del tokenizer, text_encoder, tokenizer_2, text_encoder_2
74 gc.collect()
75 torch.cuda.empty_cache()
76
77 generator = torch.Generator(device='cuda')
78 for i, generation in enumerate(queue, start=1):
79 generator.manual_seed(generation['seed'])
80 generation['latents'] = pipe(
81 prompt_embeds=generation['embeddings'][0],
82 negative_prompt_embeds=generation['embeddings'][1],
83 pooled_prompt_embeds=generation['embeddings'][2],
84 negative_pooled_prompt_embeds=generation['embeddings'][3],
85 generator=generator,
86 output_type='latent',
87 guidance_scale=generation['cfg_scale'],
88 num_inference_steps=generation['steps'],
89 height=generation['height'],
90 width=generation['width']
91 ).images
92
93
94 del pipe.unet
95 gc.collect()
96 torch.cuda.empty_cache()
97
98 # load_pipe(pipe, dir="cached_pipe")
99
100
101 pipe.upcast_vae()
102 with torch.no_grad():
103 for i, generation in enumerate(queue, start=1):
104 generation['latents'] = generation['latents'].to(next(iter(pipe.vae.post_quant_conv.parameters())).dtype)
105 image = pipe.vae.decode(generation['latents'] / pipe.vae.config.scaling_factor, return_dict=False)[0]
106 image = pipe.image_processor.postprocess(image, output_type='pil')[0]
107 image_path = f'image_{i}.png'
108 image.save(image_path)
109 print(f"Image saved at: {image_path}")
110 # save_pipe(pipe, dir="cached_pipe")
111
112
113
114 if input("Do you want to create another image? (y/n): ").lower() != 'y':
115 break
116