moralec/MagicQuill
0
1import subprocess2import shlex3subprocess.run(4 shlex.split(5 "pip install ./gradio_magicquill-0.0.1-py3-none-any.whl"6 )7)8import gradio as gr9from gradio_magicquill import MagicQuill10import random11import torch12import numpy as np13from PIL import Image, ImageOps14import base6415import io16from fastapi import FastAPI, Request17import uvicorn18from MagicQuill import folder_paths19from MagicQuill.scribble_color_edit import ScribbleColorEditModel20from gradio_client import Client, handle_file21from huggingface_hub import snapshot_download22import tempfile23import cv224import os25import requests26 27snapshot_download(repo_id="LiuZichen/MagicQuill-models", repo_type="model", local_dir="models")28# HF_TOKEN = os.environ.get("HF_TOKEN")29# The client has been made public. Welcome to duplicate our repo.30client = Client("LiuZichen/DrawNGuess")31scribbleColorEditModel = ScribbleColorEditModel()32 33def tensor_to_numpy(tensor):34 if isinstance(tensor, torch.Tensor):35 return (tensor.detach().cpu().numpy() * 255).astype(np.uint8)36 return tensor37 38def tensor_to_base64(tensor):39 tensor = tensor.squeeze(0) * 255.40 pil_image = Image.fromarray(tensor.cpu().byte().numpy())41 buffered = io.BytesIO()42 pil_image.save(buffered, format="PNG")43 img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")44 45 return img_str46 47def read_base64_image(base64_image):48 if base64_image.startswith("data:image/png;base64,"):49 base64_image = base64_image.split(",")[1]50 elif base64_image.startswith("data:image/jpeg;base64,"):51 base64_image = base64_image.split(",")[1]52 elif base64_image.startswith("data:image/webp;base64,"):53 base64_image = base64_image.split(",")[1]54 else:55 raise ValueError("Unsupported image format.")56 image_data = base64.b64decode(base64_image)57 image = Image.open(io.BytesIO(image_data))58 image = ImageOps.exif_transpose(image)59 return image60 61def create_alpha_mask(base64_image):62 """Create an alpha mask from the alpha channel of an image."""63 image = read_base64_image(base64_image)64 mask = torch.zeros((1, image.height, image.width), dtype=torch.float32, device="cpu")65 if 'A' in image.getbands():66 alpha_channel = np.array(image.getchannel('A')).astype(np.float32) / 255.067 mask[0] = 1.0 - torch.from_numpy(alpha_channel)68 return mask69 70def load_and_preprocess_image(base64_image, convert_to='RGB', has_alpha=False):71 """Load and preprocess a base64 image."""72 image = read_base64_image(base64_image)73 image = image.convert(convert_to)74 image_array = np.array(image).astype(np.float32) / 255.075 image_tensor = torch.from_numpy(image_array)[None,]76 return image_tensor77 78def load_and_resize_image(base64_image, convert_to='RGB', max_size=512):79 """Load and preprocess a base64 image, resize if necessary."""80 image = read_base64_image(base64_image)81 image = image.convert(convert_to)82 width, height = image.size83 # if min(width, height) > max_size:84 scaling_factor = max_size / min(width, height)85 new_size = (int(width * scaling_factor), int(height * scaling_factor))86 image = image.resize(new_size, Image.LANCZOS)87 image_array = np.array(image).astype(np.float32) / 255.088 image_tensor = torch.from_numpy(image_array)[None,]89 return image_tensor90 91def prepare_images_and_masks(total_mask, original_image, add_color_image, add_edge_image, remove_edge_image):92 total_mask = create_alpha_mask(total_mask)93 original_image_tensor = load_and_preprocess_image(original_image)94 if add_color_image:95 add_color_image_tensor = load_and_preprocess_image(add_color_image)96 else:97 add_color_image_tensor = original_image_tensor98 99 add_edge_mask = create_alpha_mask(add_edge_image) if add_edge_image else torch.zeros_like(total_mask)100 remove_edge_mask = create_alpha_mask(remove_edge_image) if remove_edge_image else torch.zeros_like(total_mask)101 return add_color_image_tensor, original_image_tensor, total_mask, add_edge_mask, remove_edge_mask102 103def guess_prompt_handler(original_image, add_color_image, add_edge_image):104 original_image_tensor = load_and_preprocess_image(original_image)105 106 if add_color_image:107 add_color_image_tensor = load_and_preprocess_image(add_color_image)108 else:109 add_color_image_tensor = original_image_tensor110 111 width, height = original_image_tensor.shape[1], original_image_tensor.shape[2]112 add_edge_mask = create_alpha_mask(add_edge_image) if add_edge_image else torch.zeros((1, height, width), dtype=torch.float32, device="cpu")113 114 original_image_numpy = tensor_to_numpy(original_image_tensor.squeeze(0))115 add_color_image_numpy = tensor_to_numpy(add_color_image_tensor.squeeze(0))116 add_edge_mask_numpy = tensor_to_numpy(add_edge_mask.squeeze(0).unsqueeze(-1))117 original_image_numpy = cv2.cvtColor(original_image_numpy, cv2.COLOR_RGB2BGR)118 add_color_image_numpy = cv2.cvtColor(add_color_image_numpy, cv2.COLOR_RGB2BGR)119 120 original_image_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", mode='w+b')121 add_color_image_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", mode='w+b')122 add_edge_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", mode='w+b')123 124 cv2.imwrite(original_image_file.name, original_image_numpy)125 cv2.imwrite(add_color_image_file.name, add_color_image_numpy)126 cv2.imwrite(add_edge_mask_file.name, add_edge_mask_numpy)127 128 original_image_file.close()129 add_color_image_file.close()130 add_edge_mask_file.close()131 132 res = client.predict(133 handle_file(original_image_file.name),134 handle_file(add_color_image_file.name),135 handle_file(add_edge_mask_file.name)136 )137 138 if original_image_file and os.path.exists(original_image_file.name):139 os.remove(original_image_file.name)140 if add_color_image_file and os.path.exists(add_color_image_file.name):141 os.remove(add_color_image_file.name)142 if add_edge_mask_file and os.path.exists(add_edge_mask_file.name):143 os.remove(add_edge_mask_file.name)144 145 return res146 147def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_image, remove_edge_image, positive_prompt, negative_prompt, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):148 add_color_image, original_image, total_mask, add_edge_mask, remove_edge_mask = prepare_images_and_masks(total_mask, original_image, add_color_image, add_edge_image, remove_edge_image)149 progress = None150 if fine_edge == 'disable':151 if torch.sum(remove_edge_mask).item() > 0 and torch.sum(add_edge_mask).item() == 0:152 if positive_prompt == "":153 positive_prompt = "empty scene"154 edge_strength /= 3.155 156 latent_samples, final_image, lineart_output, color_output = scribbleColorEditModel.process(157 ckpt_name,158 original_image, 159 add_color_image, 160 positive_prompt, 161 negative_prompt, 162 total_mask, 163 add_edge_mask, 164 remove_edge_mask, 165 grow_size, 166 stroke_as_edge, 167 fine_edge, 168 edge_strength, 169 color_strength, 170 inpaint_strength, 171 seed,172 steps,173 cfg,174 sampler_name,175 scheduler,176 progress177 )178 179 final_image_base64 = tensor_to_base64(final_image)180 return final_image_base64181 182def generate_image_handler(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):183 if seed == -1:184 seed = random.randint(0, 2**32 - 1)185 ms_data = x['from_frontend']186 positive_prompt = x['from_backend']['prompt']187 stroke_as_edge = "enable"188 res = generate(ckpt_name, ms_data['total_mask'], ms_data['original_image'], ms_data['add_color_image'], ms_data['add_edge_image'], ms_data['remove_edge_image'], positive_prompt, negative_prompt, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler)189 x["from_backend"]["generated_image"] = res190 return x191 192css = '''193.row {194 width: 90%;195 margin: auto;196}197'''198 199head = """200<meta http-equiv="Content-Security-Policy" content="frame-ancestors 'none'">201"""202 203with gr.Blocks(css=css, head=head) as demo:204 with gr.Row(elem_classes="row"):205 text = gr.Markdown(206 """207 # Welcome to MagicQuill!208 Click the [link](https://magicquill.art) to view our demo and tutorial. Give us a [GitHub star](https://github.com/magic-quill/magicquill) if you are interested. 209 """)210 with gr.Row(elem_classes="row"):211 ms = MagicQuill()212 with gr.Row(elem_classes="row"):213 with gr.Column():214 btn = gr.Button("Run", variant="primary")215 with gr.Column():216 with gr.Accordion("parameters", open=False):217 ckpt_name = gr.Dropdown(218 label="Base Model Name",219 choices=folder_paths.get_filename_list("checkpoints"),220 value='SD1.5/realisticVisionV60B1_v51VAE.safetensors',221 interactive=True222 )223 negative_prompt = gr.Textbox(224 label="Negative Prompt",225 value="",226 interactive=True227 )228 # stroke_as_edge = gr.Radio(229 # label="Stroke as Edge",230 # choices=['enable', 'disable'],231 # value='enable',232 # interactive=True233 # )234 fine_edge = gr.Radio(235 label="Fine Edge",236 choices=['enable', 'disable'],237 value='disable',238 interactive=True239 )240 grow_size = gr.Slider(241 label="Grow Size",242 minimum=0,243 maximum=100,244 value=15,245 step=1,246 interactive=True247 )248 edge_strength = gr.Slider(249 label="Edge Strength",250 minimum=0.0,251 maximum=5.0,252 value=0.55,253 step=0.01,254 interactive=True255 )256 color_strength = gr.Slider(257 label="Color Strength",258 minimum=0.0,259 maximum=5.0,260 value=0.55,261 step=0.01,262 interactive=True263 )264 inpaint_strength = gr.Slider(265 label="Inpaint Strength",266 minimum=0.0,267 maximum=5.0,268 value=1.0,269 step=0.01,270 interactive=True271 )272 seed = gr.Number(273 label="Seed",274 value=-1,275 precision=0,276 interactive=True277 )278 steps = gr.Slider(279 label="Steps",280 minimum=1,281 maximum=50,282 value=20,283 step=1,284 interactive=True285 )286 cfg = gr.Slider(287 label="CFG",288 minimum=0.0,289 maximum=20.0,290 value=5.0,291 step=0.1,292 interactive=True293 )294 sampler_name = gr.Dropdown(295 label="Sampler Name",296 choices=["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "ddim", "uni_pc", "uni_pc_bh2"],297 value='euler_ancestral',298 interactive=True299 )300 scheduler = gr.Dropdown(301 label="Scheduler",302 choices=["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"],303 value='karras',304 interactive=True305 )306 btn.click(generate_image_handler, inputs=[ms, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler], outputs=ms, concurrency_limit=1)307 308 with gr.Row(elem_classes="row"):309 text = gr.Markdown(310 """311 Note: This demo is governed by the license of CC BY-NC 4.0. We strongly advise users not to knowingly generate or allow others to knowingly generate harmful content, including hate speech, violence, pornography, deception, etc. (注：本演示受CC BY-NC的许可协议限制。我们强烈建议，用户不应传播及不应允许他人传播以下内容，包括但不限于仇恨言论、暴力、色情、欺诈相关的有害信息。)312 """)313 demo.queue(max_size=20, status_update_rate=0.1)314 315app = FastAPI()316 317@app.post("/magic_quill/guess_prompt")318async def guess_prompt(request: Request):319 data = await request.json()320 res = guess_prompt_handler(data['original_image'], data['add_color_image'], data['add_edge_image'])321 return res322 323@app.post("/magic_quill/process_background_img")324async def process_background_img(request: Request):325 img = await request.json()326 resized_img_tensor = load_and_resize_image(img)327 resized_img_base64 = "data:image/png;base64," + tensor_to_base64(resized_img_tensor)328 return resized_img_base64329 330app = gr.mount_gradio_app(app, demo, "/")331 332if __name__ == "__main__":333 uvicorn.run(app, host="0.0.0.0", port=7860)334 # demo.launch()