Spaces:
Sleeping
Sleeping
File size: 4,204 Bytes
2793b6f 51ff353 2793b6f 8733506 | 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 | import os
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
# --- Configuration ---
# Using the high-speed distilled FLUX model
MODEL_ID = "black-forest-labs/FLUX.1-schnell"
# Pre-defined style templates
STYLE_OPTIONS = {
"None": "{prompt}",
"Realistic": "A hyper-realistic, high-detail photograph of {prompt}, 8k resolution, cinematic lighting, sharp focus",
"Spirited Away": "Studio Ghibli style, hand-drawn animation of {prompt}, lush colors, whimsical atmosphere, cinematic",
"Cyberpunk": "Cyberpunk aesthetic, neon lighting, futuristic city vibes, {prompt}, high contrast, 8k",
"Minimalist": "Clean smooth lines, minimalist style, {prompt}, simple white background, high-quality digital art",
"Oil Painting": "Classic oil painting texture, visible brushstrokes, rich colors, {prompt}, masterpiece"
}
def authenticate_client():
"""
Retrieves the 'TextToImage' token (stored as HF_TOKEN secret)
from the Space environment.
"""
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
raise ValueError("HF_TOKEN secret not found. Please add your 'TextToImage' token to Space Secrets.")
return InferenceClient(token=hf_token)
def apply_style_to_prompt(prompt: str, style: str) -> str:
"""
Combines the user prompt with the selected style keywords.
"""
template = STYLE_OPTIONS.get(style, "{prompt}")
return template.format(prompt=prompt)
def generate_image(prompt: str, style: str):
"""
Generates the image with error handling for gated models and timeouts.
"""
try:
client = authenticate_client()
enriched_prompt = apply_style_to_prompt(prompt, style)
# Call the Inference API
# FLUX.1-schnell is optimized for quality in just 4 steps
image = client.text_to_image(
enriched_prompt,
model=MODEL_ID
)
return image
except Exception as e:
error_msg = str(e)
if "403" in error_msg or "gated" in error_msg.lower():
return gr.Error("Access Denied: Please visit the FLUX.1-schnell model card and accept the terms.")
elif "timeout" in error_msg.lower():
return gr.Error("API Timeout: The server is busy. Please try again in a moment.")
else:
return gr.Error(f"Generation Error: {error_msg}")
# --- UI Layout ---
def build_app():
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(f"# 🎨 Professional AI Image Generator")
gr.Markdown(f"Powered by `{MODEL_ID}` via Hugging Face Inference API.")
with gr.Row():
with gr.Column():
user_prompt = gr.Textbox(
label="What do you want to create?",
placeholder="Describe your vision...",
lines=3
)
style_choice = gr.Dropdown(
choices=list(STYLE_OPTIONS.keys()),
value="None",
label="Choose a Style"
)
generate_btn = gr.Button("Generate Masterpiece", variant="primary")
gr.Examples(
examples=[
["A futuristic city on Mars with glass domes", "Realistic"],
["A peaceful forest with glowing mushrooms", "Spirited Away"],
["A sleek metallic owl perched on a branch", "Minimalist"]
],
inputs=[user_prompt, style_choice]
)
with gr.Column():
output_image = gr.Image(label="Generated Result", type="pil")
# Link the button to the function
generate_btn.click(
fn=generate_image,
inputs=[user_prompt, style_choice],
outputs=output_image
)
gr.Markdown("---")
gr.Markdown("### Deployment Steps:\n1. Create a **Gradio Space**.\n2. Add your **TextToImage** token to **Settings > Secrets** as `HF_TOKEN`.\n3. Upload these files.")
return demo
if __name__ == "__main__":
app = build_app()
app.launch() |