singularityinspace/TexttoImage

1import gradio as gr2import numpy as np3import random4from diffusers import DiffusionPipeline5import torch6 7device = "cuda" if torch.cuda.is_available() else "cpu"8 9if torch.cuda.is_available():10    torch.cuda.max_memory_allocated(device=device)11    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)12    pipe.enable_xformers_memory_efficient_attention()13    pipe = pipe.to(device)14else: 15    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)16    pipe = pipe.to(device)17 18MAX_SEED = np.iinfo(np.int32).max19MAX_IMAGE_SIZE = 102420 21def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):22 23    if randomize_seed:24        seed = random.randint(0, MAX_SEED)25        26    generator = torch.Generator().manual_seed(seed)27    28    image = pipe(29        prompt = prompt, 30        negative_prompt = negative_prompt,31        guidance_scale = guidance_scale, 32        num_inference_steps = num_inference_steps, 33        width = width, 34        height = height,35        generator = generator36    ).images[0] 37    38    return image39 40examples = [41    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",42    "An astronaut riding a green horse",43    "A delicious ceviche cheesecake slice",44]45 46css="""47#col-container {48    margin: 0 auto;49    max-width: 520px;50}51"""52 53if torch.cuda.is_available():54    power_device = "GPU"55else:56    power_device = "CPU"57 58with gr.Blocks(css=css) as demo:59    60    with gr.Column(elem_id="col-container"):61        gr.Markdown(f"""62        # Text-to-Image Gradio Template63        Currently running on {power_device}.64        """)65        66        with gr.Row():67            68            prompt = gr.Text(69                label="Prompt",70                show_label=False,71                max_lines=1,72                placeholder="Enter your prompt",73                container=False,74            )75            76            run_button = gr.Button("Run", scale=0)77        78        result = gr.Image(label="Result", show_label=False)79 80        with gr.Accordion("Advanced Settings", open=False):81            82            negative_prompt = gr.Text(83                label="Negative prompt",84                max_lines=1,85                placeholder="Enter a negative prompt",86                visible=False,87            )88            89            seed = gr.Slider(90                label="Seed",91                minimum=0,92                maximum=MAX_SEED,93                step=1,94                value=0,95            )96            97            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)98            99            with gr.Row():100                101                width = gr.Slider(102                    label="Width",103                    minimum=256,104                    maximum=MAX_IMAGE_SIZE,105                    step=32,106                    value=512,107                )108                109                height = gr.Slider(110                    label="Height",111                    minimum=256,112                    maximum=MAX_IMAGE_SIZE,113                    step=32,114                    value=512,115                )116            117            with gr.Row():118                119                guidance_scale = gr.Slider(120                    label="Guidance scale",121                    minimum=0.0,122                    maximum=10.0,123                    step=0.1,124                    value=0.0,125                )126                127                num_inference_steps = gr.Slider(128                    label="Number of inference steps",129                    minimum=1,130                    maximum=12,131                    step=1,132                    value=2,133                )134        135        gr.Examples(136            examples = examples,137            inputs = [prompt]138        )139 140    run_button.click(141        fn = infer,142        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],143        outputs = [result]144    )145 146demo.queue().launch()