Update app.py

app.py

@@ -15,89 +15,9 @@ from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 import gradio as gr
 import random
 
-print("os.environ: ", os.environ)
-
 device = "cuda"
 model_path = "CompVis/stable-diffusion-v1-4"
 
-pipe = StableDiffusionInpaintingPipeline.from_pretrained(
-      model_path,
-      revision="fp16", 
-      torch_dtype=torch.float16,
-      use_auth_token=os.environ.get("access_token")).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
-    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-    image = image.resize((w, h), resample=PIL.Image.LANCZOS)
-    image = np.array(image).astype(np.float32) / 255.0
-    image = image[None].transpose(0, 3, 1, 2)
-    image = torch.from_numpy(image)
-    return 2.0 * image - 1.0
-
-def preprocess_mask(mask):
-    mask=mask.convert("L")
-    w, h = mask.size
-    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-    mask = mask.resize((w//8, h//8), resample=PIL.Image.NEAREST)
-    mask = np.array(mask).astype(np.float32) / 255.0
-    mask = np.tile(mask,(4,1,1))
-    mask = mask[None].transpose(0, 1, 2, 3)#what does this step do?
-    mask = 1 - mask #repaint white, keep black
-    mask = torch.from_numpy(mask)
-    return mask
-
-
 class StableDiffusionInpaintingPipeline(DiffusionPipeline):
     def __init__(
         self,
@@ -257,7 +177,84 @@ class StableDiffusionInpaintingPipeline(DiffusionPipeline):
             image = self.numpy_to_pil(image)
 
         return {"sample": image, "nsfw_content_detected": has_nsfw_concept}
-        
+
+pipe = StableDiffusionInpaintingPipeline.from_pretrained(
+      model_path,
+      revision="fp16", 
+      torch_dtype=torch.float16,
+      use_auth_token=os.environ.get("access_token")).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
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    image = image.resize((w, h), resample=PIL.Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return 2.0 * image - 1.0
+
+def preprocess_mask(mask):
+    mask=mask.convert("L")
+    w, h = mask.size
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    mask = mask.resize((w//8, h//8), resample=PIL.Image.NEAREST)
+    mask = np.array(mask).astype(np.float32) / 255.0
+    mask = np.tile(mask,(4,1,1))
+    mask = mask[None].transpose(0, 1, 2, 3)#what does this step do?
+    mask = 1 - mask #repaint white, keep black
+    mask = torch.from_numpy(mask)
+    return mask
+
 # 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