|
import os, torch, random |
|
from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline |
|
from kolors.models.modeling_chatglm import ChatGLMModel |
|
from kolors.models.tokenization_chatglm import ChatGLMTokenizer |
|
from diffusers import UNet2DConditionModel, AutoencoderKL |
|
from diffusers import EulerDiscreteScheduler |
|
import gradio as gr |
|
|
|
ckpt_dir = f"Kwai-Kolors/Kolors" |
|
text_encoder = ChatGLMModel.from_pretrained( |
|
f'{ckpt_dir}/text_encoder', |
|
torch_dtype=torch.float16).half() |
|
tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder') |
|
vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half() |
|
scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler") |
|
unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half() |
|
pipe = StableDiffusionXLPipeline( |
|
vae=vae, |
|
text_encoder=text_encoder, |
|
tokenizer=tokenizer, |
|
unet=unet, |
|
scheduler=scheduler, |
|
force_zeros_for_empty_prompt=False) |
|
pipe = pipe.to("cuda") |
|
pipe.enable_model_cpu_offload() |
|
|
|
def generate_image(prompt, height, width, num_inference_steps, guidance_scale): |
|
seed = random.randint(0, 18446744073709551615) |
|
image = pipe( |
|
prompt=prompt, |
|
height=height, |
|
width=width, |
|
num_inference_steps=num_inference_steps, |
|
guidance_scale=guidance_scale, |
|
num_images_per_prompt=1, |
|
generator=torch.Generator(pipe.device).manual_seed(seed) |
|
).images[0] |
|
return image, seed |
|
|
|
|
|
iface = gr.Interface( |
|
fn=generate_image, |
|
inputs=[ |
|
gr.Textbox(label="Prompt"), |
|
gr.Slider(512, 1024, 1024, step=64, label="Height"), |
|
gr.Slider(512, 1024, 1024, step=64, label="Width"), |
|
gr.Slider(20, 100, 50, step=1, label="Number of Inference Steps"), |
|
gr.Slider(1, 20, 5, step=0.5, label="Guidance Scale"), |
|
], |
|
outputs=[ |
|
gr.Image(label="Generated Image"), |
|
gr.Number(label="Seed") |
|
], |
|
title="Kolors Stable Diffusion XL Image Generator", |
|
description="Generate images using the Kolors Stable Diffusion XL model." |
|
) |
|
|
|
iface.launch() |
