ImageGeneration / app.py
Arjun Singh
new faceswapper
7f6532e
Raw
History Blame Contribute Delete
3.24 kB
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,source_image_base64):
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}"},
json={
'workflow_values': {
'negative_prompt': negative_prompt,
'positive_prompt': positive_prompt,
'seed': seed,
'source_image_fs': source_image_base64
}
}
)
result = resp.json()
print(result)
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,source_image):
seed = numpy.random.randint(0, 2**32 - 1)
source_image_base64 = None
if source_image is not None:
with open(source_image, "rb") as image_file:
source_image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
try:
return call_api_and_generate_image(negative_prompt, positive_prompt,seed,source_image_base64)
except Exception as e:
return None#, f"Error: {str(e)}"
def clear_fields():
return "", "", None # Clear prompt and outputs
with gr.Blocks(theme='freddyaboulton/test-blue') as demo:
gr.Markdown("<center><h2>Arjun's Face Swapper</h2></center>")
gr.Markdown("Hi there! Specify a man/woman in any setting below and upload your image to see yourself in that setting!")
prompt = gr.Textbox(label='Positve Prompt', lines=2, max_lines=5, placeholder = 'Describe your image here, for e.g., superman flying in the sky')
neg_prompt = gr.Textbox(label='Negative Prompt', lines=2, max_lines=10, value="Low quality, pixelated, poor anatomy, bad fingers")
source_image = gr.Image(label='Source Image for Face Swap', type="filepath")
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_box = gr.Image(type="pil", label="Final Generated Image")
#json_output = gr.Textbox(label="API Response JSON", lines=10)
submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt,source_image], outputs=[image_output_box])
clear_btn.click(fn=clear_fields,outputs=[prompt,neg_prompt,image_output_box])
demo.launch()