Update app.py

app.py

@@ -12,7 +12,43 @@ import os
 import uuid
 from datetime import datetime
 
-# Previous model initialization code remains the same...
+# Model paths
+base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
+vae_model_path = "stabilityai/sd-vae-ft-mse"
+image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
+ip_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sd15.bin", repo_type="model")
+ip_plus_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid-plusv2_sd15.bin", repo_type="model")
+
+device = "cuda"
+
+# Initialize the noise scheduler
+noise_scheduler = DDIMScheduler(
+    num_train_timesteps=1000,
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    clip_sample=False,
+    set_alpha_to_one=False,
+    steps_offset=1,
+)
+
+# Load models
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+pipe = StableDiffusionPipeline.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.float16,
+    scheduler=noise_scheduler,
+    vae=vae
+).to(device)
+
+ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
+
+# Initialize FaceAnalysis
+app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+cv2.setNumThreads(1)
 
 STYLE_PRESETS = [
     {
@@ -127,7 +163,46 @@ css = '''
 footer {display: none !important}
 '''
 
-# Generate image function remains the same...
+@spaces.GPU(enable_queue=True)
+def generate_image(images, gender, prompt, progress=gr.Progress(track_tqdm=True)):
+    if not prompt:
+        prompt = f"Professional portrait of a {gender.lower()}"
+    
+    faceid_all_embeds = []
+    first_iteration = True
+    preserve_face_structure = True
+    face_strength = 2.1
+    likeness_strength = 0.7
+
+    for image in images:
+        face = cv2.imread(image)
+        faces = app.get(face)
+        faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
+        faceid_all_embeds.append(faceid_embed)
+
+        if first_iteration and preserve_face_structure:
+            face_image = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224)
+            first_iteration = False
+
+    average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
+
+    image = ip_model_plus.generate(
+        prompt=prompt,  
+        faceid_embeds=average_embedding,
+        scale=likeness_strength, 
+        face_image=face_image, 
+        shortcut=True, 
+        s_scale=face_strength, 
+        width=512, 
+        height=912, 
+        num_inference_steps=100
+    )
+    return image
+
+def create_preset_click_handler(idx, prompt_input):
+    def handler():
+        return {"value": STYLE_PRESETS[idx]["prompt"]}
+    return handler
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_classes="container"):
@@ -149,6 +224,12 @@ with gr.Blocks(css=css) as demo:
                     type="value"
                 )
                 
+                prompt_input = gr.Textbox(
+                    label="🎨 Custom Prompt",
+                    placeholder="Describe your desired transformation in detail...",
+                    lines=3
+                )
+                
                 with gr.Column(elem_classes="preset-container"):
                     gr.Markdown("### 🎭 Magic Transformations")
                     preset_grid = []
@@ -158,16 +239,11 @@ with gr.Blocks(css=css) as demo:
                             elem_classes="preset-card"
                         )
                         preset_button.click(
-                            lambda idx=idx: (idx, STYLE_PRESETS[idx]["prompt"]),
-                            outputs=[gr.Number(visible=False), prompt_input]
+                            fn=create_preset_click_handler(idx, prompt_input),
+                            inputs=[],
+                            outputs=[prompt_input]
                         )
                         preset_grid.append(preset_button)
-
-                prompt_input = gr.Textbox(
-                    label="🎨 Custom Prompt",
-                    placeholder="Describe your desired transformation in detail...",
-                    lines=3
-                )
                 
                 generate_button = gr.Button("🚀 Generate Magic", variant="primary")