from PIL import Image from io import BytesIO import numpy as np import base64 import torch import torchvision.transforms.functional as F from S2I import Sketch2ImagePipeline class Sketch2ImageController(): def __init__(self, gr): super().__init__() self.gr = gr self.style_list = [ {"name": "Comic", "prompt": "comic {prompt} . graphic illustration, comic art, graphic novel art, vibrant, highly detailed"}, {"name": "Cinematic", "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy"}, {"name": "3D Model", "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting"}, {"name": "Anime", "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed"}, {"name": "Digital Art", "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed"}, {"name": "Photographic", "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed"}, {"name": "Pixel art", "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics"}, {"name": "Fantasy art", "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy"}, {"name": "Neonpunk", "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional"}, {"name": "Manga", "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style"}, ] self.styles = {k["name"]: k["prompt"] for k in self.style_list} self.STYLE_NAMES = list(self.styles.keys()) self.DEFAULT_STYLE_NAME = "Fantasy art" self.MAX_SEED = np.iinfo(np.int32).max # Initialize the model once here self.pipe = None self.zero_options = None def load_pipeline(self, zero_options): if self.pipe is None or zero_options != self.zero_options: self.pipe = Sketch2ImagePipeline() self.zero_options = zero_options def update_canvas(self, use_line, use_eraser): brush_size = 20 if use_eraser else 4 _color = "#ffffff" if use_eraser else "#000000" return self.gr.update(brush_radius=brush_size, brush_color=_color, interactive=True) def upload_sketch(self, file): _img = Image.open(file.name).convert("L") return self.gr.update(value=_img, source="upload", interactive=True) @staticmethod def pil_image_to_data_uri(img, format="PNG"): buffered = BytesIO() img.save(buffered, format=format) img_str = base64.b64encode(buffered.getvalue()).decode() return f"data:image/{format.lower()};base64,{img_str}" def artwork(self, options, image, prompt, prompt_template, style_name, seed, val_r, faster, model_name, type_flag): self.load_pipeline(zero_options=options) prompt = prompt_template.replace("{prompt}", prompt) if type_flag == 'live-sketch': img = Image.fromarray(np.array(image["composite"])[:, :, -1]) elif type_flag == 'upload': img = image["composite"] img = img.convert("RGB") img = img.resize((512, 512)) image_t = F.to_tensor(img) > 0.5 c_t = image_t.unsqueeze(0).cuda().float() torch.manual_seed(seed) _, _, H, W = c_t.shape noise = torch.randn((1, 4, H // 8, W // 8), device=c_t.device) with torch.no_grad(): output_image = self.pipe.generate(c_t, prompt, r=val_r, noise_map=noise, half_model=faster, model_name=model_name) output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5) if type_flag == 'live-sketch': input_uri = self.pil_image_to_data_uri(Image.fromarray(255 - np.array(img))) else: input_uri = self.pil_image_to_data_uri(img) return output_pil, self.gr.update(link=input_uri)