Update app.py

app.py

@@ -16,7 +16,67 @@ from diffusers import AutoencoderKL, DDIMScheduler, DiffusionPipeline, PNDMSched
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from tqdm.auto import tqdm
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
+import gradio as gr
+import random
 
+device = "cuda"
+model_path = "CompVis/stable-diffusion-v1-4"
+
+pipe = StableDiffusionInpaintingPipeline.from_pretrained(
+      model_path,
+      revision="fp16", 
+      torch_dtype=torch.float16,
+      use_auth_token=True).to(device)
+  
+import gdown
+def download_gdrive_url():
+  url = 'https://drive.google.com/u/0/uc?id=1PPO2MCttsmSqyB-vKh5C7SumwFKuhgyj&export=download'
+  output = 'haarcascade_frontalface_default.xml'
+  gdown.download(url, output, quiet=False)
+  
+from torch import autocast
+def inpaint(p, init_image, mask_image=None, strength=0.75, guidance_scale=7.5, generator=None, num_samples=1, n_iter=1):
+    all_images = []
+    for _ in range(n_iter):
+        with autocast("cuda"):
+            images = pipe(
+                prompt=[p] * num_samples,
+                init_image=init_image,
+                mask_image=mask_image,
+                strength=strength,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                num_inference_steps=75
+            )["sample"]
+        all_images.extend(images)
+    print(len(all_images))
+    return all_images[0]
+    
+def identify_face(user_image):
+  img = cv2.imread(user_image.name) # read the resized image in cv2
+  print(img.shape)
+  gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # convert to grayscale
+  download_gdrive_url() #download the haarcascade face recognition stuff
+  haar_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
+  faces_rect = haar_cascade.detectMultiScale(gray_img, scaleFactor=1.1, minNeighbors=9)
+  for (x, y, w, h) in faces_rect[:1]:
+    mask = np.zeros(img.shape[:2], dtype="uint8")
+    print(mask.shape)
+    cv2.rectangle(mask, (x, y), (x+w, y+h),  255, -1)
+    print(mask.shape)
+    inverted_image = cv2.bitwise_not(mask)
+    return inverted_image
+
+def sample_images(init_image, mask_image):
+  p = "4K UHD professional profile picture of a person wearing a suit for work"
+  strength=0.65
+  guidance_scale=10
+  num_samples = 1
+  n_iter = 1
+
+  generator = torch.Generator(device="cuda").manual_seed(random.randint(0, 1000000)) # change the seed to get different results
+  all_images = inpaint(p, init_image, mask_image, strength=strength, guidance_scale=guidance_scale, generator=generator, num_samples=num_samples, n_iter=n_iter)
+  return all_images
 
 def preprocess_image(image):
     w, h = image.size
@@ -199,68 +259,7 @@ class StableDiffusionInpaintingPipeline(DiffusionPipeline):
             image = self.numpy_to_pil(image)
 
         return {"sample": image, "nsfw_content_detected": has_nsfw_concept}
-  
-device = "cuda"
-model_path = "CompVis/stable-diffusion-v1-4"
-
-pipe = StableDiffusionInpaintingPipeline.from_pretrained(
-      model_path,
-      revision="fp16", 
-      torch_dtype=torch.float16,
-      use_auth_token=True).to(device)
-  
-import gdown
-def download_gdrive_url():
-  url = 'https://drive.google.com/u/0/uc?id=1PPO2MCttsmSqyB-vKh5C7SumwFKuhgyj&export=download'
-  output = 'haarcascade_frontalface_default.xml'
-  gdown.download(url, output, quiet=False)
-  
-from torch import autocast
-def inpaint(p, init_image, mask_image=None, strength=0.75, guidance_scale=7.5, generator=None, num_samples=1, n_iter=1):
-    all_images = []
-    for _ in range(n_iter):
-        with autocast("cuda"):
-            images = pipe(
-                prompt=[p] * num_samples,
-                init_image=init_image,
-                mask_image=mask_image,
-                strength=strength,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                num_inference_steps=75
-            )["sample"]
-        all_images.extend(images)
-    print(len(all_images))
-    return all_images[0]
-    
-def identify_face(user_image):
-  img = cv2.imread(user_image.name) # read the resized image in cv2
-  print(img.shape)
-  gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # convert to grayscale
-  download_gdrive_url() #download the haarcascade face recognition stuff
-  haar_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
-  faces_rect = haar_cascade.detectMultiScale(gray_img, scaleFactor=1.1, minNeighbors=9)
-  for (x, y, w, h) in faces_rect[:1]:
-    mask = np.zeros(img.shape[:2], dtype="uint8")
-    print(mask.shape)
-    cv2.rectangle(mask, (x, y), (x+w, y+h),  255, -1)
-    print(mask.shape)
-    inverted_image = cv2.bitwise_not(mask)
-    return inverted_image
-
-def sample_images(init_image, mask_image):
-  p = "4K UHD professional profile picture of a person wearing a suit for work"
-  strength=0.65
-  guidance_scale=10
-  num_samples = 1
-  n_iter = 1
-
-  generator = torch.Generator(device="cuda").manual_seed(random.randint(0, 1000000)) # change the seed to get different results
-  all_images = inpaint(p, init_image, mask_image, strength=strength, guidance_scale=guidance_scale, generator=generator, num_samples=num_samples, n_iter=n_iter)
-  return all_images
-
-import gradio as gr
-import random
+        
 # accept an image input 
 # trigger the set of functions to occur => identify face, generate mask, save the inverted face mask, sample for the inverted images
 # output the sampled images