app.py

import random
import torch
import gc
import gradio as gr
from PIL import Image
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline,\
                      StableDiffusionXLPipeline, StableDiffusionUpscalePipeline,\
                      DiffusionPipeline
from utils import *
from style import custom_css, header_html

def gen_image(prompt, negative_prompt, width, height, num_steps, mode, seed, guidance_scale, device):
    """
    Run diffusion model to generate image
    """
    device = f"cuda:{device.split('GPU')[1][1]}"
    guidance_scale = float(guidance_scale)
    generator = torch.Generator(device).manual_seed(int(seed))
    model_path = DIFFUSION_CHECKPOINTS[mode]["path"]
    Text2Image_class = globals()[DIFFUSION_CHECKPOINTS[mode]["pipeline"]]
    if DIFFUSION_CHECKPOINTS[mode]["type"] == "pretrained":
        pipeline = Text2Image_class.from_pretrained(model_path)
    else:
        pipeline = Text2Image_class.from_single_file(model_path)
    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
    try:
        pipeline = pipeline.to(device)
        image = pipeline(prompt=prompt,
                         negative_prompt=negative_prompt,
                         width=nearest_divisible_by_8(int(width)),
                         height=nearest_divisible_by_8(int(height)),
                         num_inference_steps=int(num_steps),
                         generator=generator,
                         guidance_scale=guidance_scale).images[0]
    except Exception as e:
        image = Image.open("stuffs/serverdown.jpg")
        print(e)
    del pipeline
    torch.cuda.empty_cache()
    gc.collect()
    return image

with gr.Blocks(title="TonAI Creative", theme=APP_THEME) as interface1:
    gr.HTML(header_html)
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", placeholder="Tell me what you want to generate")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Instruct the AI model that it should not include")
            with gr.Row():
                width = gr.Textbox(label="Image Width", value=768)
                height = gr.Textbox(label="Image Height", value=768)
            with gr.Row():
                seed = gr.Textbox(label="RNG Seed", value=0, scale=1)
                guidance_scale = gr.Textbox(label="CFG Scale", value=7, scale=1)
            with gr.Row():
                num_steps = gr.components.Slider(
                                minimum=5, maximum=60, value=20, step=1,
                                label="Inference Steps"
                            )
                mode=gr.Dropdown(choices=DIFFUSION_CHECKPOINTS.keys(), label="Mode",
                                 value=list(DIFFUSION_CHECKPOINTS.keys())[1])
            device_choices = display_gpu_info()
            device=gr.Dropdown(choices=device_choices, label="Device", value=device_choices[0])
            generate_btn = gr.Button("Generate")
        with gr.Column(scale=2):
            generate_btn.click(
                        fn=gen_image,
                        inputs=[prompt, negative_prompt, width, height, num_steps, mode, seed, guidance_scale, device],
                        outputs=gr.Image(label="Generated Image", format="png"),
                        concurrency_limit=10
                    )
        interface1.load(lambda: gr.update(value=random.randint(0, 999999)), None, seed)
        interface1.load(lambda: gr.update(choices=display_gpu_info(), value=display_gpu_info()[0]), None, device)

# interface = gr.TabbedInterface([interface1, iface2], ["Text-to-text", "image-to-text"])
allowed_paths=["stuffs/tonai_research_logo.png"]
interface1.queue(default_concurrency_limit=10)
interface1.launch(share=True, 
                  allowed_paths=allowed_paths,
                  max_threads=10)