Spaces:
Running on Zero
Running on Zero
| """画像専用の独立子プロセス。分離のGPU経路を利用しない。""" | |
| import hashlib | |
| import json | |
| import os | |
| from pathlib import Path | |
| import sys | |
| import warnings | |
| os.environ.setdefault('OMP_NUM_THREADS', '1') | |
| os.environ.setdefault('OPENBLAS_NUM_THREADS', '1') | |
| os.environ.setdefault('CUDA_VISIBLE_DEVICES', '') | |
| def write(folder, **values): | |
| p = folder/'progress.tmp' | |
| p.write_text(json.dumps(values), encoding='utf-8') | |
| p.replace(folder/'progress.json') | |
| def model_file(name, expected): | |
| root = Path(os.environ.get('LUNARON_IMAGE_MODEL_DIR', Path(__file__).parent/'image_models')) | |
| p = root/name | |
| if hashlib.sha256(p.read_bytes()).hexdigest() != expected: | |
| raise ValueError('MODEL_INTEGRITY') | |
| return p | |
| def open_image(p, mode): | |
| from PIL import Image, ImageOps | |
| Image.MAX_IMAGE_PIXELS = 16000000 | |
| warnings.simplefilter('error', Image.DecompressionBombWarning) | |
| im = Image.open(p) | |
| if im.format not in {'JPEG', 'PNG', 'WEBP'} or getattr(im, 'n_frames', 1) != 1: | |
| raise ValueError('INVALID_IMAGE') | |
| if im.width*im.height > 16000000: | |
| raise ValueError('IMAGE_TOO_LARGE') | |
| im.load() | |
| im = ImageOps.exif_transpose(im).convert('RGBA') | |
| if mode == 'upscale' and (max(im.size) > 1024 or im.width*im.height > 1048576): | |
| raise ValueError('UPSCALE_LIMIT') | |
| return im | |
| def background(folder, im): | |
| import numpy as np | |
| import onnxruntime as ort | |
| from PIL import Image, ImageChops | |
| path = model_file('birefnet-lite.onnx', '5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333') | |
| options = ort.SessionOptions() | |
| options.intra_op_num_threads = 1 | |
| options.inter_op_num_threads = 1 | |
| session = ort.InferenceSession(str(path), sess_options=options, providers=['CPUExecutionProvider']) | |
| rgb = Image.new('RGB', im.size, 'white') | |
| rgb.paste(im, mask=im.getchannel('A')) | |
| arr = np.asarray(rgb.resize((1024, 1024), Image.Resampling.LANCZOS), dtype=np.float32)/255.0 | |
| arr = ((arr-np.array([.485,.456,.406],dtype=np.float32))/np.array([.229,.224,.225],dtype=np.float32)).transpose(2,0,1)[None] | |
| write(folder, stage='processing') | |
| logits = session.run(None, {session.get_inputs()[0].name:arr})[0][0,0] | |
| mask = 1/(1+np.exp(-np.clip(logits,-30,30))) | |
| # 固定の確率を使い、背景がない画像をmin/max拡大で無理に切り抜かない。 | |
| mask = Image.fromarray(np.uint8(np.clip(mask*255,0,255)), 'L') | |
| im.thumbnail((2048,2048), Image.Resampling.LANCZOS) | |
| mask = mask.resize(im.size, Image.Resampling.LANCZOS) | |
| im.putalpha(ImageChops.multiply(im.getchannel('A'), mask)) | |
| write(folder, stage='saving') | |
| im.save(folder/'transparent.png', compress_level=4) | |
| white = Image.new('RGBA',im.size,'white') | |
| white.alpha_composite(im) | |
| white.convert('RGB').save(folder/'white.png', compress_level=4) | |
| return im | |
| def upscale(folder, im): | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from PIL import Image | |
| torch.set_num_threads(1) | |
| torch.set_num_interop_threads(1) | |
| # Real-ESRGAN SRVGGNetCompact。著作権・BSD-3-Clauseはimage_modelsに同梱。 | |
| class SRVGG(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.body=nn.ModuleList([nn.Conv2d(3,64,3,1,1),nn.PReLU(64)]) | |
| for _ in range(32): | |
| self.body.extend([nn.Conv2d(64,64,3,1,1),nn.PReLU(64)]) | |
| self.body.append(nn.Conv2d(64,48,3,1,1)) | |
| self.upsampler=nn.PixelShuffle(4) | |
| def forward(self,x): | |
| out=x | |
| for layer in self.body: out=layer(out) | |
| return self.upsampler(out)+F.interpolate(x,scale_factor=4,mode='nearest') | |
| p=model_file('realesr-general-x4v3.pth','8dc7edb9ac80ccdc30c3a5dca6616509367f05fbc184ad95b731f05bece96292') | |
| model=SRVGG().eval() | |
| checkpoint=torch.load(p,map_location='cpu',weights_only=True) | |
| model.load_state_dict(checkpoint['params_ema'] if 'params_ema' in checkpoint else checkpoint['params'],strict=True) | |
| rgb=Image.new('RGB',im.size,'white');rgb.paste(im,mask=im.getchannel('A')) | |
| arr=np.asarray(rgb,dtype=np.float32)/255 | |
| w,h=im.size | |
| out=Image.new('RGB',(w*2,h*2)) | |
| tiles=((w+127)//128)*((h+127)//128);done=0 | |
| with torch.inference_mode(): | |
| for y in range(0,h,128): | |
| for x in range(0,w,128): | |
| sx=max(0,x-20);sy=max(0,y-20);ex=min(w,x+148);ey=min(h,y+148) | |
| t=torch.from_numpy(arr[sy:ey,sx:ex].transpose(2,0,1).copy())[None] | |
| pred=model(t)[0].clamp(0,1).permute(1,2,0).numpy() | |
| crop=Image.fromarray(np.uint8(np.round(pred*255))) | |
| box=((x-sx)*4,(y-sy)*4,(min(w,x+128)-sx)*4,(min(h,y+128)-sy)*4) | |
| crop=crop.crop(box).resize((min(128,w-x)*2,min(128,h-y)*2),Image.Resampling.LANCZOS) | |
| out.paste(crop,(x*2,y*2));done+=1 | |
| write(folder,stage='processing',done=done,total=tiles) | |
| # 透過の輪郭は元のアルファから補間し、黒背景として潰さない。 | |
| if im.getchannel('A').getextrema() != (255,255): | |
| out=out.convert('RGBA');out.putalpha(im.getchannel('A').resize(out.size,Image.Resampling.LANCZOS)) | |
| write(folder,stage='saving') | |
| out.save(folder/'enhanced.png',compress_level=4) | |
| return out | |
| def run(folder): | |
| from PIL import Image | |
| if hasattr(os,'nice'): os.nice(10) | |
| config=json.loads((folder/'config.json').read_text(encoding='utf-8')) | |
| write(folder,stage='preparing') | |
| im=open_image(folder/'input',config['mode']) | |
| thumb=im.copy();thumb.thumbnail((480,480),Image.Resampling.LANCZOS) | |
| thumb.save(folder/'before.png') | |
| output=background(folder,im) if config['mode']=='background' else upscale(folder,im) | |
| preview=output.copy();preview.thumbnail((480,480),Image.Resampling.LANCZOS) | |
| preview.save(folder/'preview.png',compress_level=4) | |
| result={'stage':'done','width':output.width,'height':output.height} | |
| write(folder,**result) | |
| (folder/'input').unlink(missing_ok=True) | |
| return result | |
| if __name__=='__main__': | |
| folder=Path(sys.argv[1]).resolve() | |
| try: run(folder) | |
| except Exception: | |
| write(folder,stage='error',error='IMAGE_PROCESSING_FAILED') | |
| (folder/'input').unlink(missing_ok=True) | |
| sys.exit(1) | |