File size: 2,943 Bytes
2733832
56f7197
 
 
 
 
2733832
 
 
 
 
 
 
 
 
0a0bb7f
 
 
 
 
 
 
5550052
56f7197
 
2733832
56f7197
2733832
56f7197
2733832
29d0acf
 
0a0bb7f
56f7197
 
0a0bb7f
56f7197
 
 
2733832
 
56f7197
 
 
 
 
 
 
 
 
2733832
d93bc7b
56f7197
d93bc7b
29d0acf
d93bc7b
 
2733832
56f7197
2733832
 
 
56f7197
2733832
 
0a0bb7f
2733832
 
 
 
 
 
 
d93bc7b
0a0bb7f
2733832
 
 
 
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
74
75
76
77
78
79
import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import base64

# List of popular styles
STYLES = [
    "Photorealistic", "Oil Painting", "Watercolor", "Anime", 
    "Studio Ghibli", "Black and White", "Polaroid", "Sketch", 
    "3D Render", "Pixel Art", "Cyberpunk", "Steampunk", 
    "Art Nouveau", "Pop Art", "Minimalist"
]

# Default negative prompt
DEFAULT_NEGATIVE_PROMPT = """
ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, 
extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, 
cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face
"""

def enhance_prompt(client, prompt, style):
    model = client.get_model('gemini-pro')
    enhanced_prompt_request = f"Enhance the following prompt for image generation in the style of {style}. Make it more detailed and vivid, while keeping the original intent: '{prompt}'"
    
    response = model.generate_content(enhanced_prompt_request)
    
    return response.text

def generate_image(client, enhanced_prompt, style, negative_prompt):
    full_prompt = f"{enhanced_prompt}\nStyle: {style}\nNegative prompt: {negative_prompt}"
    
    response = client.models.generate_images(
        model='imagen-3.0-generate-002',
        prompt=full_prompt,
        config=types.GenerateImagesConfig(
            number_of_images=1,
        )
    )
    
    image_bytes = response.generated_images[0].image.image_bytes
    image = Image.open(BytesIO(image_bytes))
    
    # Convert PIL Image to base64 string
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    
    return f"data:image/png;base64,{img_str}"

def process_and_generate(api_key, prompt, style, negative_prompt):
    client = genai.Client(api_key=api_key)
    enhanced_prompt = enhance_prompt(client, prompt, style)
    image_url = generate_image(client, enhanced_prompt, style, negative_prompt)
    return image_url, enhanced_prompt

with gr.Blocks() as demo:
    gr.Markdown("# Google Imagen Image Generator")
    
    with gr.Row():
        with gr.Column(scale=1):
            api_key = gr.Textbox(label="Google AI API Key", type="password")
            prompt = gr.Textbox(label="Prompt")
            style = gr.Dropdown(label="Style", choices=STYLES)
            negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
            generate_btn = gr.Button("Generate Image")
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image")
            enhanced_prompt_output = gr.Textbox(label="Enhanced Prompt")
    
    generate_btn.click(
        process_and_generate,
        inputs=[api_key, prompt, style, negative_prompt],
        outputs=[image_output, enhanced_prompt_output]
    )

demo.launch()