app.py

import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline, AutoencoderTiny
from PIL import Image
import numpy as np
import random
import os
import tempfile

# ── Model Setup ──────────────────────────────────────────────────────────────
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading FLUX.1-schnell pipeline …")
try:
    taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
    pipe = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell",
        torch_dtype=dtype,
        vae=taef1,
    )
    if device == "cuda":
        pipe.enable_model_cpu_offload()
    else:
        pipe = pipe.to(device)
    torch.cuda.empty_cache() if device == "cuda" else None
    print("Pipeline ready on", device)
except Exception as e:
    print(f"Error loading pipeline: {e}")
    pipe = None

MAX_SEED = np.iinfo(np.int32).max

# ── Configuration ─────────────────────────────────────────────────────────────
STYLE_PRESETS = {
    "None": "",
    "Cinematic": "cinematic lighting, film grain, anamorphic lens, 8k uhd, ",
    "Anime": "anime illustration, vibrant colors, cel shading, detailed line art, ",
    "Photorealistic": "photorealistic, hyperdetailed, 8k, professional photography, ",
    "Digital Art": "digital painting, concept art, trending on artstation, ",
    "Minimalist": "minimalist, clean composition, soft lighting, ",
    "Cyberpunk": "cyberpunk, neon lights, futuristic city, high tech, ",
    "Fantasy": "epic fantasy, magical atmosphere, highly detailed, ",
    "Watercolor": "watercolor painting, soft brush strokes, artistic, ",
    "Sci-Fi": "sci-fi, advanced technology, sleek design, holographic, ",
}

RESOLUTIONS = {
    "Square (1024x1024)": (1024, 1024),
    "Portrait (832x1216)": (832, 1216),
    "Landscape (1216x832)": (1216, 832),
    "Wide (1344x768)": (1344, 768),
    "Tall (768x1344)": (768, 1344),
    "HD (1280x720)": (1280, 720),
    "Full HD (1920x1080)": (1920, 1080),
}

# ── Generation ────────────────────────────────────────────────────────────────
@spaces.GPU
def generate(prompt, style, resolution, seed_val, randomize_seed, guidance, steps):
    if pipe is None:
        return None, 0, None, "Error: Pipeline not loaded. Please check logs."
    
    if randomize_seed:
        seed_val = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(int(seed_val))
    width, height = RESOLUTIONS[resolution]
    full_prompt = f"{STYLE_PRESETS[style]}{prompt}" if STYLE_PRESETS[style] else prompt

    result = pipe(
        prompt=full_prompt,
        width=width,
        height=height,
        guidance_scale=guidance,
        num_inference_steps=int(steps),
        generator=generator,
        output_type="pil",
        max_sequence_length=256,
    )
    image = result.images[0]
    
    # Save for download
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    image.save(tmp.name, "PNG")
    tmp.close()
    
    if device == "cuda":
        torch.cuda.empty_cache()
    
    return image, seed_val, tmp.name, None

# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
.chain-nav {
    display: flex; align-items: center; justify-content: center;
    gap: 12px; padding: 14px; background: #1a1a2e; border-radius: 12px;
    margin-bottom: 20px; box-shadow: 0 4px 15px rgba(0,0,0,0.2);
}
.chain-nav .step {
    padding: 8px 18px; border-radius: 24px; font-size: 14px;
    font-weight: 600; text-decoration: none; transition: all 0.3s;
}
.chain-nav .step.active {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white; box-shadow: 0 2px 10px rgba(102,126,234,0.4);
}
.chain-nav .step:not(.active) {
    background: #16213e; color: #a0a0a0; border: 1px solid #0f3460;
}
.chain-nav .step:not(.active):hover {
    background: #0f3460; color: white;
}
.chain-nav .arrow { font-size: 20px; color: #e94560; }
.main-title { text-align: center; margin: 8px 0 4px; font-size: 2.2rem; }
.subtitle { text-align: center; color: #888; margin-bottom: 20px; }
"""

# ── Gradio App ────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="AI Image Forge") as demo:
    gr.HTML("""
    <div class="chain-nav">
        <span class="step active">1. AI Image Forge</span>
        <span class="arrow">&rarr;</span>
        <a href="https://huggingface.co/spaces/Crealink/color-effects-studio" target="_blank" class="step">2. Color &amp; Effects</a>
        <span class="arrow">&rarr;</span>
        <a href="https://huggingface.co/spaces/Crealink/quality-maximizer" target="_blank" class="step">3. Quality &amp; Export</a>
    </div>
    <h1 class="main-title">AI Image Forge</h1>
    <p class="subtitle">Generate stunning images with <b>FLUX.1-schnell</b> &mdash; state-of-the-art open diffusion</p>
    """)

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to create...",
                lines=3,
                show_copy_button=True,
            )
            style = gr.Dropdown(
                choices=list(STYLE_PRESETS.keys()),
                value="None",
                label="Style Preset",
            )
            resolution = gr.Dropdown(
                choices=list(RESOLUTIONS.keys()),
                value="Square (1024x1024)",
                label="Resolution",
            )
            with gr.Row():
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                seed = gr.Number(label="Seed", value=42, precision=0)
            with gr.Row():
                guidance = gr.Slider(1, 10, value=3.5, step=0.1, label="Guidance Scale")
                steps = gr.Slider(1, 20, value=4, step=1, label="Steps (FLUX-schnell)")
            generate_btn = gr.Button("Generate Image", variant="primary", size="lg")
            error_text = gr.Textbox(label="Status", interactive=False, visible=False)
            
            gr.HTML('<div style="text-align:center; margin-top:12px;"><a href="https://huggingface.co/spaces/Crealink/color-effects-studio" target="_blank" style="text-decoration:none;"><button style="padding:10px 24px; border:none; border-radius:8px; background:#e94560; color:white; font-weight:600; cursor:pointer;">Continue to Color Studio</button></a></div>')

        with gr.Column(scale=2):
            output_image = gr.Image(label="Generated Image", show_label=False, height=600)
            with gr.Row():
                output_seed = gr.Number(label="Used Seed", precision=0, interactive=False)
                download_file = gr.File(label="Download", visible=True)

    generate_btn.click(
        fn=generate,
        inputs=[prompt, style, resolution, seed, randomize_seed, guidance, steps],
        outputs=[output_image, output_seed, download_file, error_text],
    )

    gr.Examples(
        examples=[
            ["a serene Japanese garden with cherry blossoms and a koi pond"],
            ["a futuristic cityscape at sunset with flying cars and neon lights"],
            ["a portrait of a wise old wizard with a long beard and glowing staff"],
            ["an astronaut riding a horse on the moon, photorealistic"],
            ["a cozy cabin in a snowy forest with warm lights glowing from windows"],
            ["a cyberpunk street market at night with holographic signs"],
        ],
        inputs=[prompt],
        label="Example Prompts",
    )

demo.launch()