Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import argparse | |
| import os | |
| import time | |
| from os import path | |
| from safetensors.torch import load_file | |
| from huggingface_hub import hf_hub_download | |
| cache_path = path.join(path.dirname(path.abspath(__file__)), "models") | |
| os.environ["TRANSFORMERS_CACHE"] = cache_path | |
| os.environ["HF_HUB_CACHE"] = cache_path | |
| os.environ["HF_HOME"] = cache_path | |
| import gradio as gr | |
| import torch | |
| from diffusers import FluxPipeline | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| class timer: | |
| def __init__(self, method_name="timed process"): | |
| self.method = method_name | |
| def __enter__(self): | |
| self.start = time.time() | |
| print(f"{self.method} starts") | |
| def __exit__(self, exc_type, exc_val, exc_tb): | |
| end = time.time() | |
| print(f"{self.method} took {str(round(end - self.start, 2))}s") | |
| if not path.exists(cache_path): | |
| os.makedirs(cache_path, exist_ok=True) | |
| pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) | |
| pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")) | |
| pipe.fuse_lora(lora_scale=0.125) | |
| pipe.to(device="cuda", dtype=torch.bfloat16) | |
| # Define example prompts | |
| example_prompts = [ | |
| "A cyberpunk cityscape at night with neon lights reflecting in puddles, towering skyscrapers and flying cars", | |
| "An ethereal fairy with translucent iridescent wings standing in an enchanted forest with glowing mushrooms and floating light particles", | |
| "A majestic dragon soaring through stormy clouds above jagged mountain peaks as lightning strikes in the background", | |
| "A futuristic space station orbiting a vibrant nebula with multiple colorful ringed planets visible through a massive observation window", | |
| "An underwater scene of an ancient lost city with ornate temples and statues covered in bioluminescent coral and swimming sea creatures" | |
| ] | |
| # Custom CSS for neon theme | |
| css = """ | |
| .neon-container { | |
| background: linear-gradient(to right, #000428, #004e92); | |
| border-radius: 16px; | |
| box-shadow: 0 0 15px #00f3ff, 0 0 25px #00f3ff; | |
| } | |
| .neon-title { | |
| text-shadow: 0 0 5px #fff, 0 0 10px #fff, 0 0 15px #0073e6, 0 0 20px #0073e6, 0 0 25px #0073e6; | |
| color: #ffffff; | |
| font-weight: bold !important; | |
| } | |
| .neon-text { | |
| color: #00ff9d; | |
| text-shadow: 0 0 5px #00ff9d; | |
| } | |
| .neon-button { | |
| box-shadow: 0 0 5px #ff00dd, 0 0 10px #ff00dd !important; | |
| background: linear-gradient(90deg, #ff00dd, #8b00ff) !important; | |
| border: none !important; | |
| color: white !important; | |
| font-weight: bold !important; | |
| } | |
| .neon-button:hover { | |
| box-shadow: 0 0 10px #ff00dd, 0 0 20px #ff00dd !important; | |
| } | |
| .neon-input { | |
| border: 1px solid #00f3ff !important; | |
| box-shadow: 0 0 5px #00f3ff !important; | |
| } | |
| .neon-slider > div { | |
| background: linear-gradient(90deg, #00ff9d, #00f3ff) !important; | |
| } | |
| .neon-slider > div > div { | |
| background: #ff00dd !important; | |
| box-shadow: 0 0 5px #ff00dd !important; | |
| } | |
| .neon-card { | |
| background-color: rgba(0, 0, 0, 0.7) !important; | |
| border: 1px solid #00f3ff !important; | |
| box-shadow: 0 0 10px #00f3ff !important; | |
| padding: 16px !important; | |
| border-radius: 8px !important; | |
| } | |
| .neon-example { | |
| background: rgba(0, 0, 0, 0.5) !important; | |
| border: 1px solid #00ff9d !important; | |
| border-radius: 8px !important; | |
| padding: 8px !important; | |
| color: #00ff9d !important; | |
| box-shadow: 0 0 5px #00ff9d !important; | |
| margin: 4px !important; | |
| cursor: pointer !important; | |
| } | |
| .neon-example:hover { | |
| box-shadow: 0 0 10px #00ff9d, 0 0 15px #00ff9d !important; | |
| background: rgba(0, 255, 157, 0.2) !important; | |
| } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo: | |
| with gr.Blocks(elem_classes=["neon-container"]): | |
| gr.Markdown( | |
| """ | |
| <div style="text-align: center; max-width: 650px; margin: 0 auto;"> | |
| <h1 style="font-size: 3rem; font-weight: 700; margin-bottom: 1rem; display: contents;" class="neon-title">FLUX: Fast & Furious</h1> | |
| <p style="font-size: 1.2rem; margin-bottom: 1.5rem;" class="neon-text">AutoML team from ByteDance</p> | |
| </div> | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3, elem_classes=["neon-card"]): | |
| with gr.Group(): | |
| prompt = gr.Textbox( | |
| label="Your Image Description", | |
| placeholder="E.g., A serene landscape with mountains and a lake at sunset", | |
| lines=3, | |
| elem_classes=["neon-input"] | |
| ) | |
| # Examples section | |
| gr.Markdown('<p class="neon-text">Click on any example to use it:</p>') | |
| with gr.Row(): | |
| example_boxes = [gr.Button(ex[:40] + "...", elem_classes=["neon-example"]) for ex in example_prompts] | |
| # Connect example buttons to the prompt textbox | |
| for i, example_btn in enumerate(example_boxes): | |
| example_btn.click( | |
| fn=lambda x=example_prompts[i]: x, | |
| outputs=prompt | |
| ) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| with gr.Group(): | |
| with gr.Row(): | |
| height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024, | |
| elem_classes=["neon-slider"]) | |
| width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024, | |
| elem_classes=["neon-slider"]) | |
| with gr.Row(): | |
| steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8, | |
| elem_classes=["neon-slider"]) | |
| scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5, | |
| elem_classes=["neon-slider"]) | |
| seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0, | |
| elem_classes=["neon-input"]) | |
| generate_btn = gr.Button("Generate Image", variant="primary", scale=1, elem_classes=["neon-button"]) | |
| with gr.Column(scale=4, elem_classes=["neon-card"]): | |
| output = gr.Image(label="Your Generated Image") | |
| gr.Markdown( | |
| """ | |
| <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px;" class="neon-card"> | |
| <h2 style="font-size: 1.5rem; margin-bottom: 1rem;" class="neon-text">How to Use</h2> | |
| <ol style="padding-left: 1.5rem; color: #00f3ff;"> | |
| <li>Enter a detailed description of the image you want to create.</li> | |
| <li>Or click one of our exciting example prompts above!</li> | |
| <li>Adjust advanced settings if desired (tap to expand).</li> | |
| <li>Tap "Generate Image" and wait for your creation!</li> | |
| </ol> | |
| <p style="margin-top: 1rem; font-style: italic; color: #ff00dd;">Tip: Be specific in your description for best results!</p> | |
| </div> | |
| """ | |
| ) | |
| def process_image(height, width, steps, scales, prompt, seed): | |
| global pipe | |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"): | |
| return pipe( | |
| prompt=[prompt], | |
| generator=torch.Generator().manual_seed(int(seed)), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(scales), | |
| height=int(height), | |
| width=int(width), | |
| max_sequence_length=256 | |
| ).images[0] | |
| generate_btn.click( | |
| process_image, | |
| inputs=[height, width, steps, scales, prompt, seed], | |
| outputs=output | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |