| """Pinned-release Gradio app. Import spaces before torch for ZeroGPU emulation.""" |
| import os |
| from pathlib import Path |
| import spaces |
| import torch |
| import gradio as gr |
| from huggingface_hub import snapshot_download |
| from inference import prepare, chroma_coefficients, render, load_colorizer |
| import json |
|
|
| REPO=os.getenv('MODEL_ID','User-2468/mini-unet-colorizer') |
| REV=os.getenv('MODEL_REVISION','1a9eb8af2754ad2329a24cfe50d388cb559441d0') |
| SUBFOLDER=os.getenv('MODEL_SUBFOLDER','').strip('/') |
| CPU=os.getenv('ZEROGPU_CPU_TEST')=='1' |
| DEVICE='cpu' if CPU else 'cuda' |
| if CPU:torch.set_num_threads(min(4,os.cpu_count() or 1)) |
| path=Path(REPO) |
| if not path.is_dir(): |
| prefix=SUBFOLDER+'/' if SUBFOLDER else '' |
| path=Path(snapshot_download(REPO,revision=REV,allow_patterns=[prefix+'model.safetensors',prefix+'config.json'])) |
| if SUBFOLDER:path=path/SUBFOLDER |
| cfg=json.loads((path/'config.json').read_text()) |
| if cfg.get('architecture')=='DecisionColorizer': |
| from decision_model import load_decision |
| MODEL=load_decision(path,DEVICE) |
| else:MODEL=load_colorizer(path,DEVICE) |
| MODEL.eval().requires_grad_(False) |
| PARAMETERS=sum(p.numel() for p in MODEL.parameters()) |
| if PARAMETERS>=4_000_000:raise RuntimeError('Model exceeds the four-million-parameter limit.') |
| MODES=getattr(MODEL,'modes',1) |
|
|
| class PaletteView(torch.nn.Module): |
| def __init__(self,model,mode):super().__init__();self.model=model;self.mode=mode |
| def forward(self,L): |
| return self.model(L) if self.mode<0 else self.model(L,True)['all'][:,self.mode] |
|
|
| @spaces.GPU(duration=15) |
| def infer(small,radius,mode): |
| return chroma_coefficients(PaletteView(MODEL,mode),small,radius) |
|
|
| def run(image,strength,smoothing,detail,mode): |
| if image is None:raise gr.Error('Upload a photograph first.') |
| mode=int(mode) |
| if mode < -1 or mode>=MODES or (MODES==1 and mode!=-1):raise gr.Error('That colour interpretation is unavailable.') |
| try: |
| radius={'Gentle':4,'Balanced':8,'Strong':16}[smoothing] |
| size={'Standard':256,'Detailed':384,'Maximum':512}[detail] |
| light,alpha,small=prepare(image,size) |
| coefficients=infer(small,radius,mode) |
| result=render(light,alpha,coefficients,strength) |
| cache=(light,alpha,coefficients) |
| except (ValueError,RuntimeError,TypeError,KeyError) as exc: |
| raise gr.Error(str(exc)) from exc |
| return result,cache |
|
|
| def recolour(strength,cache): |
| if cache is None:return gr.skip() |
| light,alpha,coefficients=cache |
| return render(light,alpha,coefficients,strength) |
|
|
| with gr.Blocks(title='Mini Photo Colorizer',delete_cache=(3600,3600)) as demo: |
| gr.Markdown('# Give an old photograph new colour\nPlausible colour, with the original detail and proportions preserved.') |
| cache=gr.State(value=None,time_to_live=600) |
| with gr.Row(): |
| source=gr.Image(type='pil',image_mode='RGBA',label='Your photograph',sources=['upload','clipboard'],height=460) |
| result=gr.Image(type='pil',label='Colourised photograph',format='png',height=460,interactive=False) |
| with gr.Row(): |
| button=gr.Button('Colourise',variant='primary',scale=2) |
| clear=gr.ClearButton([source,result,cache],value='Clear',scale=1) |
| strength=gr.Slider(0,1.5,value=1.,step=.05,label='Colour strength',info='Adjust after colourising without another GPU request.') |
| with gr.Accordion('Fine-tune the result',open=False): |
| detail=gr.Radio(['Standard','Detailed','Maximum'],value='Standard',label='Detail',info='Higher detail can change the colours. Try Detailed if small objects are missed.') |
| smoothing=gr.Radio(['Gentle','Balanced','Strong'],value='Balanced',label='Colour smoothing') |
| palette=gr.Dropdown([('Automatic',-1)]+[(f'Alternative {i+1}',i) for i in range(MODES)] if MODES>1 else [('Automatic',-1)],value=-1,label='Colour interpretation',visible=MODES>1) |
| gr.Markdown('After changing these settings, select **Colourise** again.') |
| button.click(run,[source,strength,smoothing,detail,palette],[result,cache],api_name='colorize',concurrency_limit=1) |
| strength.release(recolour,[strength,cache],result,api_name=False,concurrency_limit=1) |
| source.change(lambda:None,outputs=cache,api_name=False,queue=False) |
| gr.Markdown('Download the result as PNG using its download button. Up to 12 megapixels; transparency is preserved.\n\nColours are interpretations, not recovered historical facts. Unusual scenes, damaged scans and tiny objects can still produce muted colours or colour bleeding. [About the model](https://huggingface.co/User-2468/mini-unet-colorizer).') |
| demo.queue(default_concurrency_limit=1,max_size=12) |
| if __name__=='__main__':demo.launch(max_file_size='30mb',ssr_mode=False) |
|
|
