Upload sdxl_sam.py

sdxl_sam.py

@@ -0,0 +1,107 @@
+# -*- coding: utf-8 -*-
+"""SDXL-SAM.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/github/pranavsrinivasa/SDXL-SAM-HEADSHOT-CREATOR/blob/main/SDXL-SAM.ipynb
+"""
+
+!pip install diffusers
+!pip install transformers
+
+from diffusers import StableDiffusionInpaintPipeline
+import torch
+
+model_id = 'stabilityai/stable-diffusion-2-inpainting'
+sd_pipeline = StableDiffusionInpaintPipeline.from_pretrained(model_id,torch_dtype = torch.float16)
+sd_pipeline = sd_pipeline.to("cuda")
+
+!pip install gradio==3.48.0
+
+!pip install -q 'git+https://github.com/facebookresearch/segment-anything.git'
+!pip install -q jupyter_bbox_widget roboflow dataclasses-json supervision
+
+!mkdir -p {HOME}/weights
+!wget -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P {HOME}/weights
+
+import os
+
+CHECKPOINT_PATH = os.path.join("{HOME}", "weights", "sam_vit_h_4b8939.pth")
+print(CHECKPOINT_PATH, "; exist:", os.path.isfile(CHECKPOINT_PATH))
+MODEL_TYPE = "vit_h"
+
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+
+sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to("cuda")
+
+predictor = SamPredictor(sam)
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+selected_pixels = []
+isInvert = 0
+
+with gr.Blocks() as genaieg:
+  selected_pixels = []
+  isInvert = 0
+  with gr.Row():
+    input_img = gr.Image(label = 'Input')
+    mask_img = gr.Image(label = "Mask")
+  with gr.Row():
+    output_img = gr.Image(label = "Ouput")
+
+  def invertmask():
+    global isInvert
+    isInvert = not(isInvert)
+
+  with gr.Row():
+    prompt_text = gr.Textbox(line = 1,label = 'Prompt')
+    submit = gr.Button('Submit')
+    radio = gr.Radio(['Invert Mask'])
+    radio.select(fn = invertmask)
+
+  def generate_mask(image, evt: gr.SelectData):
+    selected_pixels.append(evt.index)
+    predictor.set_image(image)
+    input_points = np.array(selected_pixels)
+    input_label = np.ones(input_points.shape[0])
+    mask, _, _ = predictor.predict(
+        point_coords = input_points,
+        point_labels = input_label,
+        multimask_output = False
+    )
+    if isInvert:
+      mask = np.logical_not(mask)
+    mask = Image.fromarray(mask[0,:,:])
+    return mask
+
+
+  def inpaint(img, mask, prompt):
+    img = Image.fromarray(img)
+    mask = Image.fromarray(mask)
+    img = img.resize((512,512))
+    mask = mask.resize((512,512))
+    negative_prompts = """
+    duplicate,low quality, lowest quality, bad shape,bad anatomy,
+    bad proportions, lowres,error,watermark,username,artistname,
+    signature,text,jpeg artifacts,blurry,more than one person,simple background
+    """
+    prompt_text = "Realistic professinal Headshot of a man for a profile pic" + prompt
+    output = sd_pipeline(prompt = prompt_text,
+                         image = img,
+                         negative_prompt = negative_prompts,
+                         mask_image = mask).images[0]
+
+    return output
+
+
+  input_img.select(generate_mask, [input_img],[mask_img])
+  submit.click(inpaint,
+               inputs=[input_img,mask_img,prompt_text],
+               outputs = [output_img])
+
+if __name__ == '__main__':
+  genaieg.launch(debug = True)