skyadmin/cog-webui-sd

1# inference handler for lightning ai2 3import re4import os5import logging6# import json7from pydantic import BaseModel8from typing import Any, Dict, Optional, TYPE_CHECKING9from dataclasses import dataclass10logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())11 12import lightning as L13from lightning.app.components.serve import PythonServer, Text14from lightning.app import BuildConfig15 16 17class _DefaultInputData(BaseModel):18    prompt: str19 20class _DefaultOutputData(BaseModel):21    img_data: str22    parameters: str23 24 25@dataclass26class CustomBuildConfig(BuildConfig):27    def build_commands(self):28        dir_path = "/content/"29        model_path = os.path.join(dir_path, "models/Stable-diffusion")30        # model_url = "https://huggingface.co/Hardy01/chill_watcher/resolve/main/models/Stable-diffusion/chilloutmix_NiPrunedFp32Fix.safetensors"31        model_url = "https://huggingface.co/Hardy01/chill_watcher/resolve/main/models/Stable-diffusion/chilloutmix_NiPrunedFp32Fix.safetensors"32        download_cmd = "wget -P {} {}".format(str(model_path), model_url)33        vae_url = "https://huggingface.co/Hardy01/chill_watcher/resolve/main/models/VAE/vae-ft-mse-840000-ema-pruned.ckpt"34        vae_path = os.path.join(dir_path, "models/VAE")35        down2 = "wget -P {} {}".format(str(vae_path), vae_url)36        lora_url1 = "https://huggingface.co/Hardy01/chill_watcher/resolve/main/models/Lora/koreanDollLikeness_v10.safetensors"37        lora_url2 = "https://huggingface.co/Hardy01/chill_watcher/resolve/main/models/Lora/taiwanDollLikeness_v10.safetensors"38        lora_path = os.path.join(dir_path, "models/Lora")39        down3 = "wget -P {} {}".format(str(lora_path), lora_url1)40        down4 = "wget -P {} {}".format(str(lora_path), lora_url2)41        # https://stackoverflow.com/questions/55313610/importerror-libgl-so-1-cannot-open-shared-object-file-no-such-file-or-directo42        cmd1 = "pip3 install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117"43        cmd2 = "pip3 install torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117"44        cmd_31 = "sudo apt-get update"45        cmd3 = "sudo apt-get install libgl1-mesa-glx"46        cmd4 = "sudo apt-get install libglib2.0-0"47        return [download_cmd, down2, down3, down4, cmd_31, cmd3, cmd4]48 49 50class PyTorchServer(PythonServer):51    def __init__(52        self,53        input_type: type = _DefaultInputData,54        output_type: type = _DefaultOutputData,55        **kwargs: Any,56        ):57        super().__init__(input_type=input_type, output_type=output_type, **kwargs)58 59        # Use the custom build config60        self.cloud_build_config = CustomBuildConfig()61    def setup(self):62        # need to install dependancies first to import packages63        import torch64        # Truncate version number of nightly/local build of PyTorch to not cause exceptions with CodeFormer or Safetensors65        if ".dev" in torch.__version__ or "+git" in torch.__version__:66            torch.__long_version__ = torch.__version__67            torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)68 69        from handler import initialize70        initialize()71 72    def predict(self, request):73        from modules.api.api import encode_pil_to_base6474        from modules import shared75        from modules.processing import StableDiffusionProcessingTxt2Img, process_images76        args = {77            "do_not_save_samples": True,78            "do_not_save_grid": True,79            "outpath_samples": "/content/desktop",80            "prompt": "lora:koreanDollLikeness_v15:0.66, best quality, ultra high res, (photorealistic:1.4), 1girl, beige sweater, black choker, smile, laughing, bare shoulders, solo focus, ((full body), (brown hair:1), looking at viewer",81            "negative_prompt": "paintings, sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, age spot, glans, (ugly:1.331), (duplicate:1.331), (morbid:1.21), (mutilated:1.21), (tranny:1.331), mutated hands, (poorly drawn hands:1.331), blurry, 3hands,4fingers,3arms, bad anatomy, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts,poorly drawn face,mutation,deformed",82            "sampler_name": "DPM++ SDE Karras",83            "steps": 20, # 2584            "cfg_scale": 8,85            "width": 512,86            "height": 768,87            "seed": -1,88        }89        print("&&&&&&&&&&&&&&&&&&&&&&&&",request)90        if request.prompt:91            prompt = request.prompt92            print("get prompt from request: ", prompt)93            args["prompt"] = prompt94        p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)95        processed = process_images(p)96        single_image_b64 = encode_pil_to_base64(processed.images[0]).decode('utf-8')97        return {98            "img_data": single_image_b64,99            "parameters": processed.images[0].info.get('parameters', ""),100        }101 102 103component = PyTorchServer(104   cloud_compute=L.CloudCompute('gpu', disk_size=20, idle_timeout=30)105)106# lightning run app app.py --cloud107app = L.LightningApp(component)108 109