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