Spaces:
Sleeping
Sleeping
| 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() |