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 StableDiffusionXLPipeline, LCMScheduler # from scheduling_tcd import TCDScheduler 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 = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16) pipe.to(device="cuda", dtype=torch.bfloat16) unet_state = load_file(hf_hub_download("ByteDance/Hyper-SD", "Hyper-SDXL-1step-Unet.safetensors"), device="cuda") pipe.unet.load_state_dict(unet_state) pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing ="trailing") with gr.Blocks() as demo: gr.Markdown(DESCRIPTION) with gr.Row(equal_height=False): with gr.Group(): with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0) result = gr.Gallery(label="Result", columns=NUM_IMAGES_PER_PROMPT, show_label=False) with gr.Group(): with gr.Row(visible=True): seed = gr.Slider( label="Seed", minimum=0, maximum=99999999, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(visible=True): width = gr.Slider( label="Width", minimum=256, maximum=8192, step=32, value=2048, ) height = gr.Slider( label="Height", minimum=256, maximum=8192, step=32, value=2048, ) gr.Examples( examples=examples, inputs=prompt, outputs=[result, seed], fn=generate, cache_examples=CACHE_EXAMPLES, ) def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: if randomize_seed: seed = random.randint(0, 99999999) return seed @spaces.GPU(duration=10) def process_image( height, width, prompt, seed, randomize_seed): global pipe with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"): return pipe( prompt=str,, num_inference_steps=1, guidance_scale=0., height=int(height), width=int(width), timesteps=[800], randomize_seed: bool = False, use_resolution_binning: bool = True, progress=gr.Progress(track_tqdm=True), ).images seed = int(randomize_seed_fn(seed, randomize_seed)) generator = torch.Generator().manual_seed(seed) reactive_controls = [ height, width, prompt, seed, randomize_seed] btn.click(process_image, inputs=reactive_controls, outputs=[output]) if __name__ == "__main__": demo.launch() DESCRIPTION = """ # Instant Image ### Super fast text to Image Generator. ### You may change the steps from 4 to 8, if you didn't get satisfied results. ### First Image processing takes time then images generate faster. """ if not torch.cuda.is_available(): DESCRIPTION += "\n

Running on CPU 🥶 This demo does not work on CPU.

" CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1" examples = [ "A Monkey with a happy face in the Sahara desert.", "Eiffel Tower was Made up of ICE.", "Color photo of a corgi made of transparent glass, standing on the riverside in Yosemite National Park.", "A close-up photo of a woman. She wore a blue coat with a gray dress underneath and has blue eyes.", "A litter of golden retriever puppies playing in the snow. Their heads pop out of the snow, covered in.", "an astronaut sitting in a diner, eating fries, cinematic, analog film", ]