File size: 2,696 Bytes
9b6b78e
 
 
 
 
d7d24e9
9b6b78e
 
 
1fa0c6b
9b6b78e
 
 
 
 
a30e43f
 
 
 
 
 
 
 
9b6b78e
 
 
 
 
 
b5e043c
9b6b78e
 
 
 
 
 
 
 
 
 
452f87d
a30e43f
9b6b78e
 
 
 
a30e43f
9b6b78e
 
 
6e4a082
9b6b78e
 
 
 
 
 
a30e43f
9b6b78e
 
 
a30e43f
 
 
9b6b78e
a30e43f
 
 
9b6b78e
a30e43f
 
 
9b6b78e
 
a30e43f
9b6b78e
a30e43f
9b6b78e
 
 
 
 
a30e43f
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#!/usr/bin/env python

import os
import uuid
import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline

DESCRIPTION = """# Playground v2.5"""
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

pipe = DiffusionPipeline.from_pretrained(
    "playgroundai/playground-v2.5-1024px-aesthetic",
    torch_dtype=torch.float16,
    use_safetensors=True,
    add_watermarker=False,
    variant="fp16"
)
pipe.to(device)

def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3,
    randomize_seed: bool = False,
):
    pipe.to(device)
    seed = random.randint(0, np.iinfo(np.int32).max) if randomize_seed else seed
    generator = torch.Generator().manual_seed(seed)

    images = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt if use_negative_prompt else None,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=25,
        generator=generator,
    ).images

    image_paths = [save_image(img) for img in images]
    return image_paths, seed

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(label="Prompt")
            run_button = gr.Button("Run")
        result = gr.Gallery(label="Result")
    with gr.Accordion("Advanced options", open=False):
        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
        negative_prompt = gr.Textbox(label="Negative prompt")
        seed = gr.Slider(label="Seed", minimum=0, maximum=np.iinfo(np.int32).max, step=1, value=0)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        width = gr.Slider(label="Width", minimum=256, maximum=1536, step=32, value=1024)
        height = gr.Slider(label="Height", minimum=256, maximum=1536, step=32, value=1024)
        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20, step=0.1, value=3.0)

    gr.on(
        triggers=[prompt.submit, negative_prompt.submit, run_button.click],
        fn=generate,
        inputs=[prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, randomize_seed],
        outputs=[result, seed],
        api_name="run",
    )

if __name__ == "__main__":
    demo.launch()