import spaces import gradio as gr import torch import gc, os, uuid, json from PIL import PngImagePlugin device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if os.getenv("SPACES_ZERO_GPU", None): torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False torch.backends.cuda.matmul.allow_tf32 = True torch.set_float32_matmul_precision("high") # https://pytorch.org/blog/accelerating-generative-ai-3/ def load_pipeline(): from diffusers import DiffusionPipeline pipe = DiffusionPipeline.from_pretrained( "John6666/rae-diffusion-xl-v2-sdxl-spo-pcm", custom_pipeline="lpw_stable_diffusion_xl", #custom_pipeline="nyanko7/sdxl_smoothed_energy_guidance", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, ) pipe.to("cpu") return pipe def token_auto_concat_embeds(pipe, positive, negative): max_length = pipe.tokenizer.model_max_length positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1] negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1] print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.') if max_length < positive_length or max_length < negative_length: print('Concatenated embedding.') if positive_length > negative_length: positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda") negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") else: negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda") positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") else: positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") positive_concat_embeds = [] negative_concat_embeds = [] for i in range(0, positive_ids.shape[-1], max_length): positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0]) negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0]) positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1) negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1) return positive_prompt_embeds, negative_prompt_embeds def save_image(image, metadata, output_dir): filename = str(uuid.uuid4()) + ".png" os.makedirs(output_dir, exist_ok=True) filepath = os.path.join(output_dir, filename) metadata_str = json.dumps(metadata) info = PngImagePlugin.PngInfo() info.add_text("metadata", metadata_str) image.save(filepath, "PNG", pnginfo=info) return filepath pipe = load_pipeline() @torch.inference_mode() @spaces.GPU(duration=10) def generate_image(prompt, neg_prompt, progress=gr.Progress(track_tqdm=True)): pipe.to(device) prompt += ", anime, masterpiece, best quality, very aesthetic, absurdres" neg_prompt += ", bad hands, bad feet, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], photo, deformed, disfigured, low contrast, photo, deformed, disfigured, low contrast" metadata = { "prompt": prompt, "negative_prompt": neg_prompt, "resolution": f"{1024} x {1024}", "guidance_scale": 7.0, "num_inference_steps": 28, "sampler": "Euler", } try: #positive_embeds, negative_embeds = token_auto_concat_embeds(pipe, prompt, neg_prompt) images = pipe( prompt=prompt, negative_prompt=neg_prompt, width=1024, height=1024, guidance_scale=7.0,# seg_scale=3.0, seg_applied_layers=["mid"], num_inference_steps=28, output_type="pil", clip_skip=2, ).images if images: image_paths = [ save_image(image, metadata, "./outputs") for image in images ] return image_paths except Exception as e: print(e) return [] finally: pipe.to("cpu") torch.cuda.empty_cache() gc.collect()