bugfix and additions

app.py

@@ -13,6 +13,8 @@ import imutils
 from imutils import face_utils
 import numpy as np
 from skimage import transform as tf
+import random
+import spaces
 
 base_model_path = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
 base_cache = "model-cache"
@@ -94,12 +96,26 @@ def morph_faces(image1_path, image2_path, alpha=0.5):
     cv2.imwrite(output_path, morphed_image) 
     return output_path
 
-def generate_image(face_image_1, face_image_2, prompt, negative_prompt, width, height, num_inference_steps, seed):
+def get_negative_prompt(gender):
+    if gender == "Boy":
+        return "(mascara, makeup: 1.4), (breasts, boobs, naked, nude: 1.4), (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation"
+    elif gender == "Girl":
+        return "(beard, mustache, male features: 1.4), (naked, nude: 1.4), (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation"
+    else:  # Random
+        return "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation"
+
+def generate_image(face_image_1, face_image_2, prompt, gender, width, height, num_inference_steps, seed):
     if seed is None:
         seed = int.from_bytes(os.urandom(4), "big")
 
+    if gender == "Random":
+        gender = random.choice(["Boy", "Girl"])
+
+    negative_prompt = get_negative_prompt(gender)
+    
     baby_image_path = morph_faces(face_image_1.name, face_image_2.name)
 
+    @spaces.GPU(duration=30)
     def generate_images(faceid_embeds, num_outputs=1):
         images = ip_model.generate(
             prompt=prompt,
@@ -125,8 +141,8 @@ gr_interface = gr.Interface(
     inputs=[
         gr.Image(type="file", label="First Face Image"),
         gr.Image(type="file", label="Second Face Image"),
-        gr.Textbox(value="portrait of a 6 y.o. boy, male child, 8k, HD, happy, in living room, perfect eyes, cute", label="Prompt"),
-        gr.Textbox(value="(mascara, makeup: 1.4), (breasts, boobs, naked, nude: 1.4), (beard, old, mustache, wrinkles: 1.4), (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation", label="Negative Prompt"),
+        gr.Textbox(value="portrait of a 6 y.o. child, 8k, HD, happy, in living room, perfect eyes, cute", label="Prompt"),
+        gr.Dropdown(choices=["Boy", "Girl", "Random"], value="Boy", label="Gender"),
         gr.Slider(256, 1024, value=768, step=64, label="Width"),
         gr.Slider(256, 1024, value=768, step=64, label="Height"),
         gr.Slider(1, 200, value=30, step=1, label="Number of Inference Steps"),

requirements.txt

@@ -0,0 +1,14 @@
+opencv-python-headless
+accelerate==0.24.0
+diffusers==0.23.0
+einops
+insightface==0.7.3
+onnxruntime-gpu
+torch==2.0.1
+torchvision==0.15.2
+transformers==4.35.0
+albumentations==1.1.0
+imutils
+dlib
+cmake
+git+https://github.com/tencent-ailab/IP-Adapter.git