File size: 2,794 Bytes
bd6ad77
 
 
 
 
 
 
2a4aef9
bd6ad77
 
 
 
 
 
 
 
 
5b51ea4
bd6ad77
 
 
5b51ea4
 
 
 
 
bd6ad77
 
 
 
 
 
2a4aef9
5b51ea4
 
bd6ad77
 
 
 
 
ff7fad1
18bd6bb
 
5b51ea4
18bd6bb
bd6ad77
5b51ea4
2a4aef9
bd6ad77
5b51ea4
bd6ad77
 
 
 
18bd6bb
bd6ad77
 
 
 
 
 
5b51ea4
bd6ad77
 
 
 
 
471dc6e
5b51ea4
18bd6bb
471dc6e
bd6ad77
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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):
    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': img_str
            }
        }
    )
    
    result = resp.json()
    base64_initial = result['result'][1]['data']
    image_data = base64.b64decode(base64_initial)
    image_output = Image.open(io.BytesIO(image_data))
    
    return image_output

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,source_image)
    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 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_box = gr.Image(type="pil", label="Final Generated Image")

    submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt,source_image], outputs=[image_output])
    clear_btn.click(fn=clear_fields,outputs=[prompt,neg_prompt,image_output_box])

demo.launch()