Spaces:
Running
Running
Initial commit of application
Browse files
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
|