Create app.py

app.py

@@ -0,0 +1,165 @@
+import os
+import time
+import logging
+import requests
+from io import BytesIO
+from PIL import Image
+import gradio as gr
+
+# ----------------------------
+# Logging Configuration
+# ----------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ----------------------------
+# Constants
+# ----------------------------
+HF_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
+DEFAULT_STYLES = [
+    "Realistic", "Cinematic", "Cyberpunk",
+    "Studio Lighting", "Highly Detailed", "4K"
+]
+
+# ----------------------------
+# Utility Functions
+# ----------------------------
+
+def get_hf_token():
+    """Load Hugging Face token from environment variable."""
+    token = os.getenv("HF_TOKEN")
+    if not token:
+        raise EnvironmentError("HF_TOKEN not found in environment variables.")
+    return token
+
+
+def style_prompt(user_input: str, style: str = None) -> str:
+    """Enhance prompt with selected style."""
+    if not user_input.strip():
+        raise ValueError("Prompt cannot be empty.")
+
+    if style and style != "None":
+        enhanced = f"{user_input}, {style}, ultra quality, sharp focus"
+    else:
+        enhanced = f"{user_input}, high quality"
+
+    return enhanced
+
+
+def query_hf_api(prompt, retries=3, timeout=60, seed=None):
+    """Send request to Hugging Face Inference API with retry logic."""
+    headers = {
+        "Authorization": f"Bearer {get_hf_token()}",
+        "Content-Type": "application/json"
+    }
+
+    payload = {
+        "inputs": prompt,
+        "options": {"wait_for_model": True}
+    }
+
+    if seed is not None:
+        payload["parameters"] = {"seed": seed}
+
+    for attempt in range(retries):
+        try:
+            response = requests.post(
+                HF_API_URL,
+                headers=headers,
+                json=payload,
+                timeout=timeout
+            )
+
+            if response.status_code == 200:
+                return response.content
+
+            elif response.status_code == 503:
+                logger.warning("Model loading, retrying...")
+                time.sleep(5)
+
+            elif response.status_code == 429:
+                logger.warning("Rate limit hit, retrying...")
+                time.sleep(10)
+
+            else:
+                logger.error(f"API Error: {response.text}")
+                raise RuntimeError(f"API Error: {response.text}")
+
+        except requests.exceptions.Timeout:
+            logger.warning("Timeout occurred, retrying...")
+            time.sleep(5)
+
+    raise RuntimeError("Failed after multiple retries.")
+
+
+def generate_image(prompt, style, seed):
+    """Main function for Gradio."""
+    try:
+        styled_prompt = style_prompt(prompt, style)
+        image_bytes = query_hf_api(styled_prompt, seed=seed)
+
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+
+        return image
+
+    except Exception as e:
+        logger.error(str(e))
+        return f"Error: {str(e)}"
+
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
+
+with gr.Blocks() as app:
+    gr.Markdown("# 🎨 AI Image Generator (FLUX.1-schnell)")
+    gr.Markdown("Generate high-quality images from text prompts using Hugging Face.")
+
+    with gr.Row():
+        prompt_input = gr.Textbox(
+            label="Enter your prompt",
+            placeholder="e.g., A futuristic city at sunset"
+        )
+
+    with gr.Row():
+        style_dropdown = gr.Dropdown(
+            ["None"] + DEFAULT_STYLES,
+            label="Select Style",
+            value="None"
+        )
+        seed_input = gr.Number(
+            label="Seed (optional)",
+            value=None,
+            precision=0
+        )
+
+    generate_btn = gr.Button("Generate Image")
+
+    output_image = gr.Image(label="Generated Image")
+    download_btn = gr.File(label="Download Image")
+
+    examples = gr.Examples(
+        examples=[
+            ["A dragon flying over mountains", "Cinematic", 42],
+            ["Cyberpunk city at night", "Cyberpunk", 123],
+            ["Portrait of a warrior", "Realistic", 7],
+        ],
+        inputs=[prompt_input, style_dropdown, seed_input],
+    )
+
+    def generate_and_download(prompt, style, seed):
+        image = generate_image(prompt, style, seed)
+        if isinstance(image, str):
+            return None, None
+
+        file_path = "output.png"
+        image.save(file_path)
+        return image, file_path
+
+    generate_btn.click(
+        fn=generate_and_download,
+        inputs=[prompt_input, style_dropdown, seed_input],
+        outputs=[output_image, download_btn]
+    )
+
+app.launch()