krrishD commited on
Commit
942ca44
1 Parent(s): 3d0f402

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +61 -62
app.py CHANGED
@@ -16,7 +16,67 @@ from diffusers import AutoencoderKL, DDIMScheduler, DiffusionPipeline, PNDMSched
16
  from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
17
  from tqdm.auto import tqdm
18
  from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 
 
19
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
  def preprocess_image(image):
22
  w, h = image.size
@@ -199,68 +259,7 @@ class StableDiffusionInpaintingPipeline(DiffusionPipeline):
199
  image = self.numpy_to_pil(image)
200
 
201
  return {"sample": image, "nsfw_content_detected": has_nsfw_concept}
202
-
203
- device = "cuda"
204
- model_path = "CompVis/stable-diffusion-v1-4"
205
-
206
- pipe = StableDiffusionInpaintingPipeline.from_pretrained(
207
- model_path,
208
- revision="fp16",
209
- torch_dtype=torch.float16,
210
- use_auth_token=True).to(device)
211
-
212
- import gdown
213
- def download_gdrive_url():
214
- url = 'https://drive.google.com/u/0/uc?id=1PPO2MCttsmSqyB-vKh5C7SumwFKuhgyj&export=download'
215
- output = 'haarcascade_frontalface_default.xml'
216
- gdown.download(url, output, quiet=False)
217
-
218
- from torch import autocast
219
- def inpaint(p, init_image, mask_image=None, strength=0.75, guidance_scale=7.5, generator=None, num_samples=1, n_iter=1):
220
- all_images = []
221
- for _ in range(n_iter):
222
- with autocast("cuda"):
223
- images = pipe(
224
- prompt=[p] * num_samples,
225
- init_image=init_image,
226
- mask_image=mask_image,
227
- strength=strength,
228
- guidance_scale=guidance_scale,
229
- generator=generator,
230
- num_inference_steps=75
231
- )["sample"]
232
- all_images.extend(images)
233
- print(len(all_images))
234
- return all_images[0]
235
-
236
- def identify_face(user_image):
237
- img = cv2.imread(user_image.name) # read the resized image in cv2
238
- print(img.shape)
239
- gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # convert to grayscale
240
- download_gdrive_url() #download the haarcascade face recognition stuff
241
- haar_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
242
- faces_rect = haar_cascade.detectMultiScale(gray_img, scaleFactor=1.1, minNeighbors=9)
243
- for (x, y, w, h) in faces_rect[:1]:
244
- mask = np.zeros(img.shape[:2], dtype="uint8")
245
- print(mask.shape)
246
- cv2.rectangle(mask, (x, y), (x+w, y+h), 255, -1)
247
- print(mask.shape)
248
- inverted_image = cv2.bitwise_not(mask)
249
- return inverted_image
250
-
251
- def sample_images(init_image, mask_image):
252
- p = "4K UHD professional profile picture of a person wearing a suit for work"
253
- strength=0.65
254
- guidance_scale=10
255
- num_samples = 1
256
- n_iter = 1
257
-
258
- generator = torch.Generator(device="cuda").manual_seed(random.randint(0, 1000000)) # change the seed to get different results
259
- all_images = inpaint(p, init_image, mask_image, strength=strength, guidance_scale=guidance_scale, generator=generator, num_samples=num_samples, n_iter=n_iter)
260
- return all_images
261
-
262
- import gradio as gr
263
- import random
264
  # accept an image input
265
  # trigger the set of functions to occur => identify face, generate mask, save the inverted face mask, sample for the inverted images
266
  # output the sampled images
 
16
  from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
17
  from tqdm.auto import tqdm
18
  from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
19
+ import gradio as gr
20
+ import random
21
 
22
+ device = "cuda"
23
+ model_path = "CompVis/stable-diffusion-v1-4"
24
+
25
+ pipe = StableDiffusionInpaintingPipeline.from_pretrained(
26
+ model_path,
27
+ revision="fp16",
28
+ torch_dtype=torch.float16,
29
+ use_auth_token=True).to(device)
30
+
31
+ import gdown
32
+ def download_gdrive_url():
33
+ url = 'https://drive.google.com/u/0/uc?id=1PPO2MCttsmSqyB-vKh5C7SumwFKuhgyj&export=download'
34
+ output = 'haarcascade_frontalface_default.xml'
35
+ gdown.download(url, output, quiet=False)
36
+
37
+ from torch import autocast
38
+ def inpaint(p, init_image, mask_image=None, strength=0.75, guidance_scale=7.5, generator=None, num_samples=1, n_iter=1):
39
+ all_images = []
40
+ for _ in range(n_iter):
41
+ with autocast("cuda"):
42
+ images = pipe(
43
+ prompt=[p] * num_samples,
44
+ init_image=init_image,
45
+ mask_image=mask_image,
46
+ strength=strength,
47
+ guidance_scale=guidance_scale,
48
+ generator=generator,
49
+ num_inference_steps=75
50
+ )["sample"]
51
+ all_images.extend(images)
52
+ print(len(all_images))
53
+ return all_images[0]
54
+
55
+ def identify_face(user_image):
56
+ img = cv2.imread(user_image.name) # read the resized image in cv2
57
+ print(img.shape)
58
+ gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # convert to grayscale
59
+ download_gdrive_url() #download the haarcascade face recognition stuff
60
+ haar_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
61
+ faces_rect = haar_cascade.detectMultiScale(gray_img, scaleFactor=1.1, minNeighbors=9)
62
+ for (x, y, w, h) in faces_rect[:1]:
63
+ mask = np.zeros(img.shape[:2], dtype="uint8")
64
+ print(mask.shape)
65
+ cv2.rectangle(mask, (x, y), (x+w, y+h), 255, -1)
66
+ print(mask.shape)
67
+ inverted_image = cv2.bitwise_not(mask)
68
+ return inverted_image
69
+
70
+ def sample_images(init_image, mask_image):
71
+ p = "4K UHD professional profile picture of a person wearing a suit for work"
72
+ strength=0.65
73
+ guidance_scale=10
74
+ num_samples = 1
75
+ n_iter = 1
76
+
77
+ generator = torch.Generator(device="cuda").manual_seed(random.randint(0, 1000000)) # change the seed to get different results
78
+ all_images = inpaint(p, init_image, mask_image, strength=strength, guidance_scale=guidance_scale, generator=generator, num_samples=num_samples, n_iter=n_iter)
79
+ return all_images
80
 
81
  def preprocess_image(image):
82
  w, h = image.size
 
259
  image = self.numpy_to_pil(image)
260
 
261
  return {"sample": image, "nsfw_content_detected": has_nsfw_concept}
262
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
263
  # accept an image input
264
  # trigger the set of functions to occur => identify face, generate mask, save the inverted face mask, sample for the inverted images
265
  # output the sampled images