Update app.py

app.py

@@ -1,16 +1,15 @@
+import torch
+import spaces
+from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
+from transformers import AutoFeatureExtractor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus
+from huggingface_hub import hf_hub_download
+from insightface.app import FaceAnalysis
+from insightface.utils import face_align
 import gradio as gr
-import asyncio
-import fal_client
-from PIL import Image
-import requests
-import io
-import os
-
-# Set up your Fal API key as an environment variable
-os.environ["FAL_KEY"] = "b6fa8d06-4225-4ec3-9aaf-4d01e960d899:cc6a52d0fc818c6f892b2760fd341ee4"
-fal_client.api_key = os.environ["FAL_KEY"]
-
-# Base model paths for model switching
+import cv2
+
 base_model_paths = {
     "RealisticVisionV4": "SG161222/Realistic_Vision_V4.0_noVAE",
     "RealisticVisionV6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
@@ -20,140 +19,142 @@ base_model_paths = {
     "EpicRealism": "emilianJR/epiCRealism"
 }
 
-# Updated function to include the API call to the Fal model
-async def generate_image(image_url: str, prompt: str, negative_prompt: str, model_type: str, base_model: str, seed: int, guidance_scale: float, num_inference_steps: int, width: int, height: int):
-    """
-    Submit the image generation process using the fal_client's submit method with the ip-adapter-face-id model.
-    Arguments:
-        image_url: URL of the input image to use.
-        prompt: Text prompt for generating the image.
-        negative_prompt: Text for negative prompt to avoid unwanted characteristics in the output.
-        model_type: Model type to use.
-        base_model: Base model to use for image generation.
-        seed: Seed for random generation.
-        guidance_scale: CFG scale for how closely the model sticks to the prompt.
-        num_inference_steps: Number of inference steps.
-        width: Width of the generated image.
-        height: Height of the generated image.
-    Returns:
-        The URL of the generated image along with other attributes like file size, dimensions, etc., or None if failed.
-    """
-    try:
-        handler = fal_client.submit(
-            "fal-ai/ip-adapter-face-id",
-            arguments={
-                "model_type": model_type,
-                "prompt": prompt,
-                "face_image_url": image_url,
-                "negative_prompt": negative_prompt,
-                "seed": seed,
-                "guidance_scale": guidance_scale,
-                "num_inference_steps": num_inference_steps,
-                "num_samples": 1,  # Adjusted to 1 sample
-                "width": width,
-                "height": height,
-                "face_id_det_size": 640,
-                "base_1_5_model_repo": base_model_paths[base_model],  # Base model
-                "base_sdxl_model_repo": "SG161222/RealVisXL_V3.0",
-                "face_images_data_url": None
-            },
-        )
-        # Retrieve the result synchronously
-        result = handler.get()
-
-        if "image" in result and "url" in result["image"]:
-            return result["image"]  # Return the full image information dictionary
-        else:
-            return None
-    except Exception as e:
-        print(f"Error generating image: {e}")
-        return None
-
-# Fetch the image from the given URL
-def fetch_image_from_url(url: str) -> Image.Image:
-    response = requests.get(url)
-    return Image.open(io.BytesIO(response.content))
-
-# Process input images and handle the image generation
-async def process_inputs(image: Image.Image, prompt: str, negative_prompt: str, model_type: str, base_model: str, seed: int, guidance_scale: float, num_inference_steps: int, width: int, height: int):
-    image_url = await upload_image_to_server(image)
-    
-    if not image_url:
-        return None
-    
-    image_info = await generate_image(image_url, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, width, height)
-    
-    if image_info and "url" in image_info:
-        return fetch_image_from_url(image_info["url"]), image_info  # Return both the image and the metadata
-    
-    return None, None
 
-# Upload image to server
-async def upload_image_to_server(image: Image.Image) -> str:
-    byte_arr = io.BytesIO()
-    image.save(byte_arr, format='PNG')
-    byte_arr.seek(0)
+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")
+
+safety_model_id = "CompVis/stable-diffusion-safety-checker"
+safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
+safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+
+device = "cuda"
 
-    with open("temp_image.png", "wb") as f:
-        f.write(byte_arr.getvalue())
+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,
+)
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
 
-    try:
-        upload_url = await fal_client.upload_file_async("temp_image.png")
-        return upload_url
-    except Exception as e:
-        print(f"Error uploading image: {e}")
-        return ""
+def load_model(base_model_path):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        base_model_path,
+        torch_dtype=torch.float16,
+        scheduler=noise_scheduler,
+        vae=vae,
+        feature_extractor=safety_feature_extractor,
+        safety_checker=None  # <--- Disable safety checker
+    ).to(device)
+    return pipe
+
+ip_model = None
+ip_model_plus = None
+
+app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+cv2.setNumThreads(1)
+
+@spaces.GPU(enable_queue=True)
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, base_model, num_inference_steps, guidance_scale, width, height, progress=gr.Progress(track_tqdm=True)):
+    global ip_model, ip_model_plus
+    base_model_path = base_model_paths[base_model]
+    pipe = load_model(base_model_path)
+    ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+    ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
+
+    faceid_all_embeds = []
+    first_iteration = True
+    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) # you can also segment the face
+            first_iteration = False
+            
+    average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
+    
+    total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"
+    
+    if(not preserve_face_structure):
+        print("Generating normal")
+        image = ip_model.generate(
+            prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
+            scale=likeness_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+        )
+    else:
+        print("Generating plus")
+        image = ip_model_plus.generate(
+            prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
+            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+        )
+    print(image)
+    return image
 
-# Change style between Photorealistic and Stylized
 def change_style(style):
     if style == "Photorealistic":
-        return gr.update(value=True), gr.update(value=1.3), gr.update(value=1.0)
+        return(gr.update(value=True), gr.update(value=1.3), gr.update(value=1.0))
     else:
-        return gr.update(value=True), gr.update(value=0.1), gr.update(value=0.8)
+        return(gr.update(value=True), gr.update(value=0.1), gr.update(value=0.8))
 
-# Gradio Interface
-def gradio_interface(image, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, width, height):
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    
-    result_image, image_info = loop.run_until_complete(
-        process_inputs(image, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, width, height)
-    )
-    
-    if result_image:
-        metadata = f"File Name: {image_info['file_name']}\nFile Size: {image_info['file_size']} bytes\nDimensions: {image_info['width']}x{image_info['height']} px\nSeed: {image_info.get('seed', 'N/A')}"
-        return result_image, metadata
-    return None, "Error generating image"
-
-# Main Gradio App
-with gr.Blocks() as demo:
-    gr.Markdown("## Image Generation with Fal API and Gradio")
-    
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(label="Upload Image", type="pil")
-            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate", lines=2)
-            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="Describe elements to avoid", lines=2)
-            style = gr.Radio(label="Generation type", choices=["Photorealistic", "Stylized"], value="Photorealistic")
-            model_type = gr.Dropdown(label="Model Type", choices=["1_5-v1", "SDXL-v2-plus", "1_5-auraface-v1"], value="SDXL-v2-plus")
-            base_model = gr.Dropdown(label="Base Model", choices=list(base_model_paths.keys()), value="RealisticVisionV4")
-            seed_input = gr.Slider(label="Seed", value=42, minimum=0, maximum=1000, step=1)
-            guidance_scale_input = gr.Slider(label="Guidance Scale", value=7.5, minimum=1.0, maximum=20.0, step=0.1)
-            num_inference_steps_input = gr.Slider(label="Inference Steps", value=50, minimum=10, maximum=100, step=1)
-            width_input = gr.Slider(label="Width", value=1024, minimum=512, maximum=1024, step=64)
-            height_input = gr.Slider(label="Height", value=1024, minimum=512, maximum=1024, step=64)
-            generate_button = gr.Button("Generate Image")
-        
-        with gr.Column():
-            generated_image = gr.Image(label="Generated Image")
-            metadata_output = gr.Textbox(label="Image Metadata", interactive=False, lines=6)
+def swap_to_gallery(images):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
 
-    generate_button.click(
-        fn=gradio_interface,
-        inputs=[image_input, prompt_input, negative_prompt_input, model_type, base_model, seed_input, guidance_scale_input, num_inference_steps_input, width_input, height_input],
-        outputs=[generated_image, metadata_output]
-    )
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
 
-    style.change(fn=change_style, inputs=style, outputs=[model_type, guidance_scale_input, num_inference_steps_input])
+css = '''
+h1{margin-bottom: 0 !important}
+footer{display:none !important}
+'''
 
-demo.launch()
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("")
+    gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            files = gr.Files(
+                        label="Drag 1 or more photos of your face",
+                        file_types=["image"]
+                    )
+            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=125)
+            with gr.Column(visible=False) as clear_button:
+                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+            prompt = gr.Textbox(label="Prompt",
+                       info="Try something like 'a photo of a man/woman/person'",
+                       placeholder="A photo of a [man/woman/person]...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality")
+            style = gr.Radio(label="Generation type", info="For stylized try prompts like 'a watercolor painting of a woman'", choices=["Photorealistic", "Stylized"], value="Photorealistic")
+            base_model = gr.Dropdown(label="Base Model", choices=list(base_model_paths.keys()), value="Realistic_Vision_V4.0_noVAE")
+            submit = gr.Button("Submit")
+            with gr.Accordion(open=False, label="Advanced Options"):
+                preserve = gr.Checkbox(label="Preserve Face Structure", info="Higher quality, less versatility (the face structure of your first photo will be preserved). Unchecking this will use the v1 model.", value=True)
+                face_strength = gr.Slider(label="Face Structure strength", info="Only applied if preserve face structure is checked", value=1.3, step=0.1, minimum=0, maximum=3)
+                likeness_strength = gr.Slider(label="Face Embed strength", value=1.0, step=0.1, minimum=0, maximum=5)
+                nfaa_negative_prompts = gr.Textbox(label="Appended Negative Prompts", info="Negative prompts to steer generations towards safe for all audiences outputs", value="naked, bikini, skimpy, scanty, bare skin, lingerie, swimsuit, exposed, see-through")
+                num_inference_steps = gr.Slider(label="Number of Inference Steps", value=30, step=1, minimum=10, maximum=100)
+                guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.1, minimum=1, maximum=20)
+                width = gr.Slider(label="Width", value=512, step=64, minimum=256, maximum=1024)
+                height = gr.Slider(label="Height", value=512, step=64, minimum=256, maximum=1024)
+        with gr.Column():
+            gallery = gr.Gallery(label="Generated Images")
+        style.change(fn=change_style,
+                    inputs=style,
+                    outputs=[preserve, face_strength, likeness_strength])
+        files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
+        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
+        submit.click(fn=generate_image,
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, base_model, num_inference_steps, guidance_scale, width, height],
+                    outputs=gallery)
+    
+    gr.Markdown("")
+    
+demo.launch()