File size: 4,698 Bytes
226a7b7
b45db27
 
 
 
 
b9ea548
b45db27
5850fbf
 
b45db27
 
 
 
 
 
 
 
 
 
 
 
 
c9c788f
b45db27
 
 
 
 
 
 
 
 
 
 
 
 
c9c788f
b45db27
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c9c788f
 
b45db27
 
 
c9c788f
b45db27
 
 
 
 
 
 
 
 
 
c9c788f
 
 
 
 
 
 
 
b45db27
 
 
 
 
 
 
 
 
 
 
 
 
c9c788f
b45db27
5664ff3
b45db27
 
 
 
 
 
c9c788f
b45db27
 
 
b9ea548
 
 
 
 
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import os
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
import gradio as gr
import random
import tqdm
from pyngrok import ngrok

NGROK_API_KEY = os.getenv("NGROK_API_KEY")

# Enable TQDM progress tracking
tqdm.monitor_interval = 0

# Load the diffusion pipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
    "kayfahaarukku/UrangDiffusion-1.0", 
    torch_dtype=torch.float16, 
    custom_pipeline="lpw_stable_diffusion_xl",
    use_safetensors=True, 
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

# Function to generate an image
def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
    pipe.to('cuda')  # Move the model to GPU when the function is called
    
    if randomize_seed:
        seed = random.randint(0, 99999999)
    if use_defaults:
        prompt = f"{prompt}, masterpiece, best quality"
        negative_prompt = f"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, {negative_prompt}"
    generator = torch.manual_seed(seed)
    
    def callback(step, timestep, latents):
        progress(step / num_inference_steps)
        return
    
    width, height = map(int, resolution.split('x'))
    image = pipe(
        prompt, 
        negative_prompt=negative_prompt,
        width=width,
        height=height, 
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        callback=callback,
        callback_steps=1
    ).images[0]

    torch.cuda.empty_cache()
    pipe.to('cpu')  # Move the model back to CPU after generation

    return image, seed

# Define Gradio interface
def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
    image, seed = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
    return image, seed, gr.update(value=seed)

def reset_inputs():
    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='1024x1024'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=False)

with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
    gr.HTML(
        "<h1>UrangDiffusion 1.0 Demo</h1>"
        )
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
            negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
            use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
            resolution_input = gr.Radio(
                choices=[
                    "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
                    "1344x768", "768x1344", "1536x640", "640x1536"
                ],
                label="Resolution",
                value="896x1152"
            )
            guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
            num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
            seed_input = gr.Slider(minimum=0, maximum=99999999, step=1, label="Seed", value=0, interactive=True)
            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=False)
            generate_button = gr.Button("Generate")
            reset_button = gr.Button("Reset")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Generated Image")

    generate_button.click(
        interface_fn,
        inputs=[
            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
        ],
        outputs=[output_image, seed_input]
    )
    
    reset_button.click(
        reset_inputs,
        inputs=[],
        outputs=[
            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
        ]
    )

# Set up ngrok
public_url = ngrok.connect(api_key=NGROK_API_KEY)
print(f"Public URL: {public_url}")

demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)