File size: 4,486 Bytes
55a3c9a 1d411fa 55a3c9a c10fea8 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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 == 'url-sketch':
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)
|