Update app.py

app.py

@@ -10,19 +10,33 @@ import os
 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 for model switching
 base_model_paths = {
-    "Realistic Vision V4": "SG161222/Realistic_Vision_V4.0_noVAE",
-    "Realistic Vision V6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "RealisticVisionV4": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "RealisticVisionV6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
     "Deliberate": "Yntec/Deliberate",
-    "Deliberate V2": "Yntec/Deliberate2",
-    "Dreamshaper 8": "Lykon/dreamshaper-8",
-    "Epic Realism": "emilianJR/epiCRealism"
+    "DeliberateV2": "Yntec/Deliberate2",
+    "Dreamshaper8": "Lykon/dreamshaper-8",
+    "EpicRealism": "emilianJR/epiCRealism"
 }
 
-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):
+# 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(
@@ -35,11 +49,13 @@ async def generate_image(image_url: str, prompt: str, negative_prompt: str, mode
                 "seed": seed,
                 "guidance_scale": guidance_scale,
                 "num_inference_steps": num_inference_steps,
-                "num_samples": num_samples,
+                "num_samples": 1,  # Adjusted to 1 sample
                 "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
+                "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
@@ -53,45 +69,34 @@ async def generate_image(image_url: str, prompt: str, negative_prompt: str, mode
         print(f"Error generating image: {e}")
         return None
 
+# Fetch the image from the given URL
 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
+# 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
     
-    # 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)
+    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:
-    """
-    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
@@ -99,84 +104,56 @@ async def upload_image_to_server(image: Image.Image) -> str:
         print(f"Error uploading image: {e}")
         return ""
 
+# Change style between Photorealistic and Stylized
 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)
     else:
         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.
-    """
+# 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)
     
-    # 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))
+    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:
-        # 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
+# 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
             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")
-
-            # Seed input
-            seed_input = gr.Number(label="Seed", value=42, precision=0)
-
-            # Guidance scale slider
-            guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.1, minimum=1, maximum=20)
-
-            # Inference steps slider
-            num_inference_steps = gr.Slider(label="Number of Inference Steps", value=50, step=1, minimum=10, maximum=100)
-
-            # 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
+            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():
-            # 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],
+        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]
     )
 
-# Launch the Gradio interface
+    style.change(fn=change_style, inputs=style, outputs=[model_type, guidance_scale_input, num_inference_steps_input])
+
 demo.launch()