File size: 4,128 Bytes
2733832
69bfb06
56f7197
4ccafa7
56f7197
cc9c490
2733832
 
 
 
 
 
 
 
 
0a0bb7f
 
 
 
 
 
 
4ccafa7
 
4dc95a1
 
 
 
 
 
 
 
 
 
 
 
 
 
2733832
56f7197
2733832
4dc95a1
 
 
 
 
 
 
 
 
 
 
2733832
cc9c490
 
0a0bb7f
cc9c490
 
 
 
 
68ee1c7
cc9c490
 
 
 
 
 
 
 
 
 
4ccafa7
cc9c490
 
 
 
 
 
 
2733832
cc9c490
 
4ccafa7
 
cc9c490
d93bc7b
 
2733832
cc9c490
2733832
 
 
cc9c490
 
2733832
 
0a0bb7f
2733832
 
 
 
 
 
 
d93bc7b
cc9c490
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import gradio as gr
import google.generativeai as genai
from PIL import Image
import io
import base64
import requests

# 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(prompt, style):
    model = genai.GenerativeModel("gemini-2.0-flash-lite")
    enhanced_prompt_request = f"""
    Task: Enhance the following prompt for image generation.
    Style: {style}
    Original prompt: '{prompt}'
    
    Instructions:
    1. Expand the prompt to be more detailed and vivid.
    2. Incorporate elements of the specified style.
    3. Maintain the original intent of the prompt.
    4. Provide ONLY the enhanced prompt, without any explanations or options.
    5. Keep the enhanced prompt concise, ideally under 100 words.

    Enhanced prompt:
    """
    
    response = model.generate_content(enhanced_prompt_request)
    
    # Extract only the enhanced prompt, removing any potential explanations or extra text
    enhanced_prompt = response.text.strip()
    
    # If the response starts with "Enhanced prompt:" or similar, remove it
    prefixes_to_remove = ["Enhanced prompt:", "Here's the enhanced prompt:", "The enhanced prompt is:"]
    for prefix in prefixes_to_remove:
        if enhanced_prompt.lower().startswith(prefix.lower()):
            enhanced_prompt = enhanced_prompt[len(prefix):].strip()
    
    return enhanced_prompt


def generate_image(stability_api_key, enhanced_prompt, style, negative_prompt):
    url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
    
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {stability_api_key}"
    }
    
    payload = {
        "model": "sd3.5-large-turbo",
        "prompt": f"{enhanced_prompt}, Style: {style}",
        "negative_prompt": negative_prompt,
        "width": 1024,
        "height": 1024,
        "num_images": 1,
        "steps": 30,
        "cfg_scale": 7.5
    }
    
    response = requests.post(url, json=payload, headers=headers)
    
    if response.status_code == 200:
        image_data = response.json()['images'][0]['image_base64']
        return f"data:image/png;base64,{image_data}"
    else:
        return f"Image generation failed: {response.text}"

def process_and_generate(google_api_key, stability_api_key, prompt, style, negative_prompt):
    genai.configure(api_key=google_api_key)
    
    enhanced_prompt = enhance_prompt(prompt, style)
    image_url = generate_image(stability_api_key, enhanced_prompt, style, negative_prompt)
    return image_url, enhanced_prompt

with gr.Blocks() as demo:
    gr.Markdown("# Stability AI Image Generator with Google Gemini Prompt Enhancement")
    
    with gr.Row():
        with gr.Column(scale=1):
            google_api_key = gr.Textbox(label="Google AI API Key", type="password")
            stability_api_key = gr.Textbox(label="Stability 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=[google_api_key, stability_api_key, prompt, style, negative_prompt],
        outputs=[image_output, enhanced_prompt_output]
    )

demo.launch()