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("""
1. AI Image Forge → 2. Color & Effects → 3. Quality & Export

AI Image Forge

Generate stunning images with FLUX.1-schnell — state-of-the-art open diffusion

""") 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('
') 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()