Update app.py

app.py

@@ -1,160 +1,182 @@
-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 cv2
-
+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"]
+
+# Model choices (base models)
 base_model_paths = {
-    "RealisticVisionV4": "SG161222/Realistic_Vision_V4.0_noVAE",
-    "RealisticVisionV6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "Realistic Vision V4": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "Realistic Vision V6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
     "Deliberate": "Yntec/Deliberate",
-    "DeliberateV2": "Yntec/Deliberate2",
-    "Dreamshaper8": "Lykon/dreamshaper-8",
-    "EpicRealism": "emilianJR/epiCRealism"
+    "Deliberate V2": "Yntec/Deliberate2",
+    "Dreamshaper 8": "Lykon/dreamshaper-8",
+    "Epic Realism": "emilianJR/epiCRealism"
 }
 
-
-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"
-
-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)
-
-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)
+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, num_samples: int, width: int, height: int):
+    """
+    Submit the image generation process using the fal_client's submit method with the ip-adapter-face-id model.
+    """
+    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": num_samples,
+                "width": width,
+                "height": height,
+                "base_1_5_model_repo": base_model_paths[base_model],  # Base model selected by user
+                "base_sdxl_model_repo": "SG161222/RealVisXL_V3.0",  # SDXL model as default
+            },
+        )
+        # 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
+
+def fetch_image_from_url(url: str) -> Image.Image:
+    """
+    Download the image from the given URL and return it as a PIL Image.
+    """
+    response = requests.get(url)
+    return Image.open(io.BytesIO(response.content))
+
+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, num_samples: int, width: int, height: int):
+    """
+    Asynchronous function to handle image upload, prompt inputs and generate the final image.
+    """
+    # Upload the image and get a valid URL
+    image_url = await upload_image_to_server(image)
     
-    total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"
+    if not image_url:
+        return None
     
-    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
+    # Run the image generation
+    image_info = await generate_image(image_url, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, num_samples, 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
+
+async def upload_image_to_server(image: Image.Image) -> str:
+    """
+    Upload an image to the fal_client and return the uploaded image URL.
+    """
+    # Convert PIL image to byte stream for upload
+    byte_arr = io.BytesIO()
+    image.save(byte_arr, format='PNG')
+    byte_arr.seek(0)
+
+    # Convert BytesIO to a file-like object that fal_client can handle
+    with open("temp_image.png", "wb") as f:
+        f.write(byte_arr.getvalue())
+
+    # Upload the image using fal_client's asynchronous method
+    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 change_style(style):
+    """
+    Changes the style for 'Photorealistic' or 'Stylized' generation type.
+    """
     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)
+
+def gradio_interface(image, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, num_samples, width, height):
+    """
+    Wrapper function to run asynchronous code in a synchronous environment like Gradio.
+    """
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    
+    # Execute the async process_inputs function
+    result_image, image_info = loop.run_until_complete(process_inputs(image, prompt, negative_prompt, model_type, base_model, seed, guidance_scale, num_inference_steps, num_samples, width, height))
+    if result_image:
+        # Display both the image and metadata
+        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"
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Image Generation with Fal API and Gradio")
+    
+    with gr.Row():
+        with gr.Column():
+            # Image input
+            image_input = gr.Image(label="Upload Image", type="pil")
+            
+            # Textbox for prompt
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate", lines=2)
+            
+            # Textbox for negative prompt
+            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="Describe elements to avoid", lines=2)
+            
+            # Radio buttons for model type (Photorealistic or Stylized)
+            style = gr.Radio(label="Generation type", choices=["Photorealistic", "Stylized"], value="Photorealistic")
+            
+            # Dropdown for selecting the base model
+            base_model = gr.Dropdown(label="Base Model", choices=list(base_model_paths.keys()), value="Realistic Vision V4")
 
-def swap_to_gallery(images):
-    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+            # Seed input
+            seed_input = gr.Number(label="Seed", value=42, precision=0)
 
-def remove_back_to_files():
-    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            # Guidance scale slider
+            guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.1, minimum=1, maximum=20)
 
-css = '''
-h1{margin-bottom: 0 !important}
-footer{display:none !important}
-'''
+            # Inference steps slider
+            num_inference_steps = gr.Slider(label="Number of Inference Steps", value=50, step=1, minimum=10, maximum=100)
 
-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)
+            # Samples slider
+            num_samples = gr.Slider(label="Number of Samples", value=4, step=1, minimum=1, maximum=10)
+
+            # Image dimensions sliders
+            width = gr.Slider(label="Width", value=1024, step=64, minimum=256, maximum=1024)
+            height = gr.Slider(label="Height", value=1024, step=64, minimum=256, maximum=1024)
+
+            # Button to trigger image generation
+            generate_button = gr.Button("Generate Image")
+        
         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()
+            # Display generated image and metadata
+            generated_image = gr.Image(label="Generated Image")
+            metadata_output = gr.Textbox(label="Image Metadata", interactive=False, lines=6)
+
+    # Style change functionality
+    style.change(fn=change_style, inputs=style, outputs=[guidance_scale, num_samples, width])
+
+    # Define the interaction between inputs and output
+    generate_button.click(
+        fn=gradio_interface,
+        inputs=[image_input, prompt_input, negative_prompt_input, style, base_model, seed_input, guidance_scale, num_inference_steps, num_samples, width, height],
+        outputs=[generated_image, metadata_output]
+    )
+
+# Launch the Gradio interface
+demo.launch()