new face swapper

app.py

@@ -15,10 +15,15 @@ import numpy
 API_KEY = os.getenv('BASETEN_API_KEY')
 API_url = os.getenv('BASETEN_COMFY_MODEL_URL')
 
-def call_api_and_generate_image(negative_prompt, positive_prompt,seed):
+def call_api_and_generate_image(negative_prompt, positive_prompt,seed, source_image):
     if not API_KEY:
         return None, "Error: API key not found in environment variables. Please set BASETEN_API_KEY."
 
+    # Convert the source image to base64
+    buffered = io.BytesIO()
+    source_image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
     resp = requests.post(
         API_url,
         headers={"Authorization": f"Api-Key {API_KEY}"},
@@ -26,41 +31,48 @@ def call_api_and_generate_image(negative_prompt, positive_prompt,seed):
             'workflow_values': {
                 'negative_prompt': negative_prompt,
                 'positive_prompt': positive_prompt,
-                'seed': seed
+                'seed': seed,
+                'source_image': img_str
             }
         }
     )
     
     result = resp.json()
-    base64_data = result['result'][0]['data']
-    image_data = base64.b64decode(base64_data)
-    image = Image.open(io.BytesIO(image_data))
-    return image, json.dumps(result, indent=2)
+    base64_initial = result['result'][0]['data']
+    initial_image_data = base64.b64decode(base64_initial)
+    initial_image = Image.open(io.BytesIO(initial_image_data))
+    
+    base64_final = result['result'][1]['data']
+    final_image_data = base64.b64decode(base64_final)
+    final_image = Image.open(io.BytesIO(final_image_data))
+    
+    return initial_image,final_image
 
-def generate_image(positive_prompt, negative_prompt):
+def generate_image(positive_prompt, negative_prompt,source_image):
     seed = numpy.random.randint(0, 2**32 - 1)
     try:
-        return call_api_and_generate_image(negative_prompt, positive_prompt,seed)
+        return call_api_and_generate_image(negative_prompt, positive_prompt,seed,source_image)
     except Exception as e:
         return None, f"Error: {str(e)}"
 
 def clear_fields():
-    return "", "", ""  # Clear prompt and output, but not the API key 
+    return "", "", None, None  # Clear prompt and outputs
 
 with gr.Blocks(theme='freddyaboulton/test-blue') as demo:
     gr.Markdown("<center><h2>Arjun's Image Generator</h2></center>")
     gr.Markdown("Hi there! I'm an AI assistant tasked with generating images.")
     prompt = gr.Textbox(label='Positve Prompt', lines=2, max_lines=5, placeholder = 'Describe your image here.')
     neg_prompt = gr.Textbox(label='Negative Prompt', lines=2, max_lines=10, value="Low quality, pixelated")
-
+    source_image = gr.Image(label='Source Image for Face Swap', type="pil")
     with gr.Group():
         with gr.Row():
                 submit_btn = gr.Button(value="Submit", elem_id="generate_button", variant="primary", size="sm")
                 clear_btn = gr.ClearButton(value="Clear Question and AI Response", elem_id="clear_button", variant="secondary", size="sm")
         gr.Markdown("<center><h3>AI Response</h3></center>")
-        image_output = gr.Image(type="pil", label="Generated Image")
-        json_output = gr.Textbox(label="API Response JSON", lines=10)
-    submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt], outputs=[image_output,json_output])
-    clear_btn.click(fn=clear_fields,outputs=[prompt,image_output,json_output, neg_prompt])
+        initial_image_output = gr.Image(type="pil", label="Initial Generated Image")
+        final_image_output = gr.Image(type="pil", label="Final Image After Face Swap")
+
+    submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt,source_image], outputs=[initial_image_output,final_image_output])
+    clear_btn.click(fn=clear_fields,outputs=[prompt,neg_prompt,initial_image_output,final_image_output])
 
 demo.launch()

app_simple_img_generator.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import requests
+import base64
+from PIL import Image
+import io
+import json
+import os
+import numpy
+#from dotenv import load_dotenv
+
+#dotenv_path = '/.env'
+#load_dotenv(dotenv_path)
+
+# Get the API key from the environment variable
+API_KEY = os.getenv('BASETEN_API_KEY')
+API_url = os.getenv('BASETEN_COMFY_MODEL_URL')
+
+def call_api_and_generate_image(negative_prompt, positive_prompt,seed):
+    if not API_KEY:
+        return None, "Error: API key not found in environment variables. Please set BASETEN_API_KEY."
+
+    resp = requests.post(
+        API_url,
+        headers={"Authorization": f"Api-Key {API_KEY}"},
+        json={
+            'workflow_values': {
+                'negative_prompt': negative_prompt,
+                'positive_prompt': positive_prompt,
+                'seed': seed
+            }
+        }
+    )
+    
+    result = resp.json()
+    base64_data = result['result'][0]['data']
+    image_data = base64.b64decode(base64_data)
+    image = Image.open(io.BytesIO(image_data))
+    return image, json.dumps(result, indent=2)
+
+def generate_image(positive_prompt, negative_prompt):
+    seed = numpy.random.randint(0, 2**32 - 1)
+    try:
+        return call_api_and_generate_image(negative_prompt, positive_prompt,seed)
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def clear_fields():
+    return "", "", ""  # Clear prompt and output, but not the API key 
+
+with gr.Blocks(theme='freddyaboulton/test-blue') as demo:
+    gr.Markdown("<center><h2>Arjun's Image Generator</h2></center>")
+    gr.Markdown("Hi there! I'm an AI assistant tasked with generating images.")
+    prompt = gr.Textbox(label='Positve Prompt', lines=2, max_lines=5, placeholder = 'Describe your image here.')
+    neg_prompt = gr.Textbox(label='Negative Prompt', lines=2, max_lines=10, value="Low quality, pixelated")
+
+    with gr.Group():
+        with gr.Row():
+                submit_btn = gr.Button(value="Submit", elem_id="generate_button", variant="primary", size="sm")
+                clear_btn = gr.ClearButton(value="Clear Question and AI Response", elem_id="clear_button", variant="secondary", size="sm")
+        gr.Markdown("<center><h3>AI Response</h3></center>")
+        image_output = gr.Image(type="pil", label="Generated Image")
+        json_output = gr.Textbox(label="API Response JSON", lines=10)
+    submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt], outputs=[image_output,json_output])
+    clear_btn.click(fn=clear_fields,outputs=[prompt,image_output,json_output, neg_prompt])
+
+demo.launch()