HemaAM commited on
Commit
5aebe57
1 Parent(s): fa9f87b

Initial commit of application

Browse files
Files changed (2) hide show
  1. app.py +53 -0
  2. utils.py +204 -0
app.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from utils import *
3
+
4
+ styles_mapping = {
5
+ "Illustration Style": '<illustration-style>', "Line Art":'<line-art>',
6
+ "Hitokomoru Style":'<hitokomoru-style-nao>', "Marc Allante": '<Marc_Allante>',
7
+ "Midjourney":'<midjourney-style>', "Hanfu Anime": '<hanfu-anime-style>',
8
+ "Birb Style": '<birb-style>'
9
+ }
10
+ with gr.Blocks() as interface:
11
+ #gr.HTML(value=HTML_TEMPLATE, show_label=False)
12
+ with gr.Row():
13
+ text_input = gr.Textbox(
14
+ label="Enter your prompt",
15
+ placeholder="Cats fighting on the road.....",
16
+ )
17
+ concept_dropdown = gr.Dropdown(
18
+ label="Select a Concept",
19
+ choices=["Illustration Style", "Line Art", "Hitokomoru Style", "Marc Allante", "Midjourney", "Hanfu Anime", "Birb Style"],
20
+ value='Marc Allante'
21
+ )
22
+
23
+ method_dropdown = gr.Dropdown(
24
+ label="Select Guidance Type",
25
+ choices=["Edge", "Contrast", "Sharpness", "Blue", "Brightness"],
26
+ value='Contrast'
27
+ )
28
+
29
+ seed_slider = gr.Slider(
30
+ label="Random Seed",
31
+ minimum=0,
32
+ maximum=2000,
33
+ step=1,
34
+ value=42
35
+ )
36
+ inputs = [text_input, concept_dropdown, method_dropdown, seed_slider]
37
+
38
+ with gr.Row():
39
+ outputs = gr.Gallery(
40
+ label="Generative Images", show_label=True,
41
+ columns=[2], rows=[1], object_fit="contain"
42
+ )
43
+
44
+ with gr.Row():
45
+ button = gr.Button("Generate Image")
46
+ button.click(show_image, inputs=inputs, outputs=outputs)
47
+
48
+ with gr.Row():
49
+ gr.Examples(examples=get_examples(), inputs=inputs, outputs=outputs, fn=show_image, cache_examples=True)
50
+
51
+
52
+ if __name__ == "__main__":
53
+ interface.launch(enable_queue=True)
utils.py ADDED
@@ -0,0 +1,204 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import PIL
2
+ import torch
3
+ import numpy as np
4
+ from PIL import Image
5
+ from tqdm import tqdm
6
+ import torch.nn.functional as F
7
+ import torchvision.transforms as T
8
+ from diffusers import LMSDiscreteScheduler, DiffusionPipeline
9
+
10
+ # configurations
11
+ torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
12
+ height, width = 512, 512
13
+ guidance_scale = 8
14
+ loss_scale = 200
15
+ num_inference_steps = 50
16
+
17
+
18
+ model_path = "CompVis/stable-diffusion-v1-4"
19
+ sd_pipeline = DiffusionPipeline.from_pretrained(
20
+ model_path,
21
+ low_cpu_mem_usage = True,
22
+ torch_dtype=torch.float32
23
+ ).to(torch_device)
24
+
25
+
26
+ sd_pipeline.load_textual_inversion("sd-concepts-library/illustration-style")
27
+ sd_pipeline.load_textual_inversion("sd-concepts-library/line-art")
28
+ sd_pipeline.load_textual_inversion("sd-concepts-library/hitokomoru-style-nao")
29
+ sd_pipeline.load_textual_inversion("sd-concepts-library/style-of-marc-allante")
30
+ sd_pipeline.load_textual_inversion("sd-concepts-library/midjourney-style")
31
+ sd_pipeline.load_textual_inversion("sd-concepts-library/hanfu-anime-style")
32
+ sd_pipeline.load_textual_inversion("sd-concepts-library/birb-style")
33
+
34
+
35
+ styles_mapping = {
36
+ "Illustration Style": '<illustration-style>', "Line Art":'<line-art>',
37
+ "Hitokomoru Style":'<hitokomoru-style-nao>', "Marc Allante": '<Marc_Allante>',
38
+ "Midjourney":'<midjourney-style>', "Hanfu Anime": '<hanfu-anime-style>',
39
+ "Birb Style": '<birb-style>'
40
+ }
41
+
42
+ # Define seeds for all the styles
43
+ seed_list = [11, 56, 110, 65, 5, 29, 47]
44
+
45
+ # Loss Function based on Edge Detection
46
+ def edge_detection(image):
47
+ channels = image.shape[1]
48
+
49
+ # Define the kernels for Edge Detection
50
+ ed_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
51
+ ed_y = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
52
+
53
+ # Replicate the Edge detection kernels for each channel
54
+ ed_x = ed_x.repeat(channels, 1, 1, 1).to(image.device)
55
+ ed_y = ed_y.repeat(channels, 1, 1, 1).to(image.device)
56
+
57
+ # ed_x = ed_x.to(torch.float16)
58
+ # ed_y = ed_y.to(torch.float16)
59
+
60
+ # Convolve the image with the Edge detection kernels
61
+ conv_ed_x = F.conv2d(image, ed_x, padding=1, groups=channels)
62
+ conv_ed_y = F.conv2d(image, ed_y, padding=1, groups=channels)
63
+
64
+ # Combine the x and y gradients after convolution
65
+ ed_value = torch.sqrt(conv_ed_x**2 + conv_ed_y**2)
66
+
67
+ return ed_value
68
+
69
+ def edge_loss(image):
70
+ ed_value = edge_detection(image)
71
+ ed_capped = (ed_value > 0.5).to(torch.float32)
72
+ return F.mse_loss(ed_value, ed_capped)
73
+
74
+ def compute_loss(original_image, loss_type):
75
+
76
+ if loss_type == 'blue':
77
+ # blue loss
78
+ # [:,2] -> all images in batch, only the blue channel
79
+ error = torch.abs(original_image[:,2] - 0.9).mean()
80
+ elif loss_type == 'edge':
81
+ # edge loss
82
+ error = edge_loss(original_image)
83
+ elif loss_type == 'contrast':
84
+ # RGB to Gray loss
85
+ transformed_image = T.functional.adjust_contrast(original_image, contrast_factor = 2)
86
+ error = torch.abs(transformed_image - original_image).mean()
87
+ elif loss_type == 'brightness':
88
+ # brightnesss loss
89
+ transformed_image = T.functional.adjust_brightness(original_image, brightness_factor = 2)
90
+ error = torch.abs(transformed_image - original_image).mean()
91
+ elif loss_type == 'sharpness':
92
+ # sharpness loss
93
+ transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor = 2)
94
+ error = torch.abs(transformed_image - original_image).mean()
95
+ elif loss_type == 'saturation':
96
+ # saturation loss
97
+ transformed_image = T.functional.adjust_saturation(original_image, saturation_factor = 10)
98
+ error = torch.abs(transformed_image - original_image).mean()
99
+ else:
100
+ print("error. Loss not defined")
101
+
102
+ return error
103
+
104
+
105
+ def get_examples():
106
+ examples = [
107
+ ['A bird sitting on a tree', 'Midjourney', 'edge', 5],
108
+ ['Cats fighting on the road', 'Marc Allante', 'brightness', 65],
109
+ ['A mouse with the head of a puppy', 'Hitokomoru Style', 'contrast', 110],
110
+ ['A woman with a smiling face in front of an Italian Pizza', 'Hanfu Anime', 'brightness', 29],
111
+ ['A campfire (oil on canvas)', 'Birb Style', 'blue', 47],
112
+ ]
113
+ return(examples)
114
+
115
+
116
+ def latents_to_pil(latents):
117
+ # bath of latents -> list of images
118
+ latents = (1 / 0.18215) * latents
119
+ with torch.no_grad():
120
+ image = sd_pipeline.vae.decode(latents).sample
121
+ image = (image / 2 + 0.5).clamp(0, 1) # 0 to 1
122
+ image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
123
+ image = (image * 255).round().astype("uint8")
124
+ return Image.fromarray(image[0])
125
+
126
+
127
+ def show_image(prompt, concept, guidance_type, seed):
128
+
129
+ prompt = f"{prompt} in the style of {styles_mapping[concept]}"
130
+ styled_image_without_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=False))
131
+ styled_image_with_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=True))
132
+ return([styled_image_without_loss, styled_image_with_loss])
133
+
134
+
135
+ def generate_image(seed, prompt, loss_type, loss_flag=False):
136
+
137
+ generator = torch.manual_seed(seed)
138
+ batch_size = 1
139
+
140
+ # scheduler
141
+ scheduler = LMSDiscreteScheduler(beta_start = 0.00085, beta_end = 0.012, beta_schedule = "scaled_linear", num_train_timesteps = 1000)
142
+ scheduler.set_timesteps(num_inference_steps)
143
+ scheduler.timesteps = scheduler.timesteps.to(torch.float32)
144
+
145
+ # text embeddings of the prompt
146
+ text_input = sd_pipeline.tokenizer(prompt, padding='max_length', max_length = sd_pipeline.tokenizer.model_max_length, truncation= True, return_tensors="pt")
147
+ input_ids = text_input.input_ids.to(torch_device)
148
+
149
+ with torch.no_grad():
150
+ text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0]
151
+
152
+ max_length = text_input.input_ids.shape[-1]
153
+ uncond_input = sd_pipeline.tokenizer(
154
+ [""] * batch_size, padding="max_length", max_length= max_length, return_tensors="pt"
155
+ )
156
+
157
+ with torch.no_grad():
158
+ uncond_embeddings = sd_pipeline.text_encoder(uncond_input.input_ids.to(torch_device))[0]
159
+
160
+ text_embeddings = torch.cat([uncond_embeddings,text_embeddings]) # shape: 2,77,768
161
+
162
+ # random latent
163
+ latents = torch.randn(
164
+ (batch_size, sd_pipeline.unet.config.in_channels, height// 8, width //8),
165
+ generator = generator,
166
+ ) .to(torch.float32)
167
+
168
+
169
+ latents = latents.to(torch_device)
170
+ latents = latents * scheduler.init_noise_sigma
171
+
172
+ for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)):
173
+
174
+ latent_model_input = torch.cat([latents] * 2)
175
+ sigma = scheduler.sigmas[i]
176
+ latent_model_input = scheduler.scale_model_input(latent_model_input, t)
177
+
178
+ with torch.no_grad():
179
+ noise_pred = sd_pipeline.unet(latent_model_input.to(torch.float32), t, encoder_hidden_states=text_embeddings)["sample"]
180
+
181
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
182
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
183
+
184
+ if loss_flag and i%5 == 0:
185
+
186
+ latents = latents.detach().requires_grad_()
187
+ # the following line alone does not work, it requires change to reduce step only once
188
+ # hence commenting it out
189
+ #latents_x0 = scheduler.step(noise_pred,t, latents).pred_original_sample
190
+ latents_x0 = latents - sigma * noise_pred
191
+
192
+ # use vae to decode the image
193
+ denoised_images = sd_pipeline.vae.decode((1/ 0.18215) * latents_x0).sample / 2 + 0.5 # range(0,1)
194
+
195
+ loss = compute_loss(denoised_images, loss_type) * loss_scale
196
+ #loss = loss.to(torch.float16)
197
+ print(f"{i} loss {loss}")
198
+
199
+ cond_grad = torch.autograd.grad(loss, latents)[0]
200
+ latents = latents.detach() - cond_grad * sigma**2
201
+
202
+ latents = scheduler.step(noise_pred,t, latents).prev_sample
203
+
204
+ return latents