Spaces:
Running
Running
File size: 3,820 Bytes
fc3288f ef04dc7 fc3288f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import numpy as np
import os
import cv2
from PIL import Image, ImageDraw
import insightface
from insightface.app import FaceAnalysis
# Diffusion
model_base = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16, use_safetensors=True, safety_checker=None,)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
lora_model_path = "./loralucy3/checkpoint-95000"
pipe.unet.load_attn_procs(lora_model_path)
pipe.to("cuda")
# Insightface model
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))
def face_swap(src_img, dest_img):
src_img = Image.open('./images/' + src_img + '.JPG')
# Convert to RGB
src_img = src_img.convert(mode='RGB')
dest_img = dest_img.convert(mode='RGB')
# Convert to array
src_img_arr = np.asarray(src_img)
dest_img_arr = np.asarray(dest_img)
# Face detection
src_faces = app.get(src_img_arr)
dest_faces = app.get(dest_img_arr)
# Initialize swapper
swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False)
# Swap face
res = dest_img_arr.copy()
for face in dest_faces:
res = swapper.get(res, face, src_faces[0], paste_back=True)
# Convert to PIL image
final_image = Image.fromarray(np.uint8(res)).convert('RGB')
return final_image
def greet(description,color,features,occasion,type_,face):
# Parse input
prompt = 'white background '
description = 'description:' + description.replace(' ', '-')
color = ' color:' + ','.join(color)
features = ' features:' + ','.join(features)
occasion = ' occasion:' + ','.join(occasion)
type_ = ' type:' + ','.join(type_)
prompt += description + color + features + occasion + type_
print('prompt:',prompt)
image = pipe(
prompt,
negative_prompt='deformed face,bad anatomy',
width=312,
height=512,
num_inference_steps=100,
guidance_scale=7.5,
cross_attention_kwargs={"scale": 1.0}
).images[0]
if(face != 'Normal'):
image = face_swap(face, image)
return image
iface = gr.Interface(fn=greet,
inputs=[gr.Textbox(label='Description'),
gr.Dropdown(label='Color',choices=['Beige','Black','Blue','Brown','Green','Grey','Orange','Pink','Purple','Red','White','Yellow'],multiselect=True),
gr.Dropdown(label='Features',choices=['3/4-sleeve','Babydoll','Closed-Back','Corset','Crochet','Cutouts','Draped','Floral','Gloves','Halter','Lace','Long','Long-Sleeve','Midi','No-Slit','Off-The-Shoulder','One-Shoulder','Open-Back','Pockets','Print','Puff-Sleeve','Ruched','Satin','Sequins','Shimmer','Short','Short-Sleeve','Side-Slit','Square-Neck','Strapless','Sweetheart-Neck','Tight','V-Neck','Velvet','Wrap'],multiselect=True),
gr.Dropdown(label='Occasion',choices=['Homecoming','Casual','Wedding-Guest','Festival','Sorority','Day','Vacation','Summer','Pool-Party','Birthday','Date-Night','Party','Holiday','Winter-Formal','Valentines-Day','Prom','Graduation'],multiselect=True),
gr.Dropdown(label='Type',choices=['Mini-Dresses','Midi-Dresses','Maxi-Dresses','Two-Piece-Sets','Rompers','Jeans','Jumpsuits','Pants','Tops','Jumpers/Cardigans','Skirts','Shorts','Bodysuits','Swimwear'],multiselect=True),
gr.Dropdown(label='Face',choices=['Normal','Cat','Lisa','Mila'], value='Cat'),
],
outputs=gr.Image(type="pil", label="Final Image", width=312, height=512))
iface.launch() |