File size: 2,339 Bytes
4358e59
9c70ba0
3fec1fb
4358e59
 
3976ed9
f4ff201
3fec1fb
 
70e3d12
3fec1fb
 
 
 
70e3d12
3fec1fb
 
 
24176cc
7e9a760
 
9c70ba0
548031b
4358e59
98b5af6
7e9a760
809a5ae
9c70ba0
9f0ced2
4358e59
9c70ba0
3976ed9
4358e59
24176cc
4358e59
809a5ae
f9aa80f
3fec1fb
4358e59
 
98b5af6
ca74145
a4ca4f9
ca74145
4358e59
 
 
 
686e886
 
4358e59
 
 
75f237b
3fec1fb
0a14984
3fec1fb
70e3d12
a4ca4f9
548031b
02c54be
24176cc
75f237b
5f1159f
a4ca4f9
769a722
a4ca4f9
24176cc
b47c647
4358e59
6fed0f7
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
from diffusers import AutoPipelineForText2Image, StableDiffusionImg2ImgPipeline
from PIL import Image
import gradio as gr
import random
import torch
import math

css = """
.btn-green {
  background-image: linear-gradient(to bottom right, #6dd178, #00a613) !important;
  border-color: #22c55e !important;
  color: #166534 !important;
}
.btn-green:hover {
  background-image: linear-gradient(to bottom right, #6dd178, #6dd178) !important;
}
"""

def generate(prompt, samp_steps, seed, strength, progress=gr.Progress(track_tqdm=True)):
    if seed < 0:
        seed = random.randint(1,999999)
    image = txt2img(
        prompt,
        num_inference_steps=1,
        guidance_scale=0.0,
        generator=torch.manual_seed(seed),
    ).images[0]
    upscaled_image = image.resize((1024,1024), 1)
    final_image = img2img(
        prompt,
        upscaled_image,
        num_inference_steps=math.ceil(samp_steps/strength),
        guidance_scale=5,
        strength=strength,
        generator=torch.manual_seed(seed),
    ).images[0]
    return [final_image], seed
        
def set_base_models():
    txt2img = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype = torch.float16,
        variant = "fp16"
    )
    txt2img.to("cuda")
    img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
        "Lykon/dreamshaper-8",
        torch_dtype = torch.float16,
        variant = "fp16",
        safety_checker=None
    )
    img2img.to("cuda")
    return txt2img, img2img

with gr.Blocks(css=css) as demo:
    with gr.Column():
        prompt = gr.Textbox(label="Prompt")
        submit_btn = gr.Button("Generate", elem_classes="btn-green")
        
        with gr.Row():
            sampling_steps = gr.Slider(1, 10, value=5, step=1, label="Sampling steps")
            strength = gr.Slider(0, 1, value=0.75, step=0.05, label="Refiner strength")
            seed = gr.Number(label="Seed", value=-1, minimum=-1, precision=0)
            lastSeed = gr.Number(label="Last Seed", value=-1, interactive=False)
            
        gallery = gr.Gallery(show_label=False, preview=True, container=False, height=1100)
        
    submit_btn.click(generate, [prompt, sampling_steps, seed, strength], [gallery, lastSeed], queue=True)
    
txt2img, img2img = set_base_models()
demo.launch(debug=True)