File size: 2,156 Bytes
f03bfaf |
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 |
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
# Gradio interface
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() |