app.py

import random
import os

import gradio as gr

from style_template import styles
from novita_client import NovitaClient, MakePhotoLoRA

from PIL import Image
import base64
from io import BytesIO

# global variable
MAX_SEED = 2 ** 32 - 1
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "Photographic (Default)"


def get_noviata_client(novita_key):
    client = NovitaClient(novita_key, os.getenv('NOVITA_API_URI', None))
    return client


get_local_storage = """
    function() {
      globalThis.setStorage = (key, value)=>{
        localStorage.setItem(key, JSON.stringify(value))
      }
       globalThis.getStorage = (key, value)=>{
        return JSON.parse(localStorage.getItem(key))
      }

       const novita_key =  getStorage('novita_key')
       return [novita_key];
      }
    """


sdxl_models = ["albedobaseXL_v04_130099.safetensors",
               "altxl_v60_146691.safetensors",
               "animeArtDiffusionXL_alpha2_91872.safetensors",
               "animeArtDiffusionXL_alpha3_93120.safetensors",
               "animeIllustDiffusion_v04_117809.safetensors",
               "animagineXLV3_v30_231047.safetensors",
               "breakdomainxl_V05g_124265.safetensors",
               "brixlAMustInYour_v40Dagobah_145992.safetensors",
               "cinemaxAlphaSDXLCinema_alpha1_107473.safetensors",
               "cineroXLPhotomatic_v12aPHENO_137703.safetensors",
               "clearhungAnimeXL_v10_117716.safetensors",
               "copaxTimelessxlSDXL1_colorfulV2_100729.safetensors",
               "counterfeitxl__98184.safetensors",
               "counterfeitxl_v10_108721.safetensors",
               "crystalClearXL_ccxl_97637.safetensors",
               "dreamshaperXL09Alpha_alpha2Xl10_91562.safetensors",
               "dynavisionXLAllInOneStylized_alpha036FP16Bakedvae_99980.safetensors",
               "dynavisionXLAllInOneStylized_beta0411Bakedvae_109970.safetensors",
               "dynavisionXLAllInOneStylized_release0534bakedvae_129001.safetensors",
               "fenrisxl_145_134980.safetensors",
               "foddaxlPhotorealism_v45_122788.safetensors",
               "formulaxl_v10_104889.safetensors",
               "juggernautXL_version2_113240.safetensors",
               "juggernautXL_version5_126522.safetensors",
               "juggernautXL_v8Rundiffusion_227002.safetensors",
               "kohakuXL_alpha7_111843.safetensors",
               "LahMysteriousSDXL_v40_122478.safetensors",
               "leosamsHelloworldSDXLModel_helloworldSDXL10_112178.safetensors",
               "leosamsHelloworldSDXL_helloworldSDXL50_268813.safetensors",
               "mbbxlUltimate_v10RC_94686.safetensors",
               "moefusionSDXL_v10_114018.safetensors",
               "nightvisionXLPhotorealisticPortrait_beta0681Bakedvae_108833.safetensors",
               "nightvisionXLPhotorealisticPortrait_beta0702Bakedvae_113098.safetensors",
               "nightvisionXLPhotorealisticPortrait_release0770Bakedvae_154525.safetensors",
               "novaPrimeXL_v10_107899.safetensors",
               "pixelwave_v10_117722.safetensors",
               "protovisionXLHighFidelity3D_beta0520Bakedvae_106612.safetensors",
               "protovisionXLHighFidelity3D_release0620Bakedvae_131308.safetensors",
               "protovisionXLHighFidelity3D_release0630Bakedvae_154359.safetensors",
               "realismEngineSDXL_v05b_131513.safetensors",
               "realismEngineSDXL_v10_136287.safetensors",
               "realisticStockPhoto_v10_115618.safetensors",
               "RealitiesEdgeXL_4_122673.safetensors",
               "realvisxlV20_v20Bakedvae_129156.safetensors",
               "riotDiffusionXL_v20_139293.safetensors",
               "roxl_v10_109354.safetensors",
               "sd_xl_base_0.9.safetensors",
               "sd_xl_base_1.0.safetensors",
               "sdxlNijiSpecial_sdxlNijiSE_115638.safetensors",
               "sdxlNijiV3_sdxlNijiV3_104571.safetensors",
               "sdxlNijiV51_sdxlNijiV51_112807.safetensors",
               "sd_xl_refiner_1.0.safetensors",
               "sdxlUnstableDiffusers_v8HEAVENSWRATH_133813.safetensors",
               "sdXL_v10Refiner_91495.safetensors",
               "sdxlYamersAnimeUltra_yamersAnimeV3_121537.safetensors",
               "shikianimexl_v10_93788.safetensors",
               "stable-diffusion-xl-1.0-inpainting-0.1.safetensors",
               "stable-diffusion-xl-base-1.0.safetensors",
               "theTalosProject_v10_117893.safetensors",
               "thinkdiffusionxl_v10_145931.safetensors",
               "protovisionXLHighFidelity3D_releaseV660Bakedvae_207131.safetensors",
               "voidnoisecorexl_r1486_150780.safetensors",
               "wlopArienwlopstylexl_v10_101973.safetensors",
               "wlopSTYLEXL_v2_126171.safetensors",
               "xl13AsmodeusSFWNSFW_v22BakedVAE_111954.safetensors",
               "xxmix9realisticsdxl_v10_123235.safetensors",
               "zavychromaxl_b2_103298.safetensors",
               "zavychromaxl_v21_129006.safetensors"]
sdxl_loras = ['sdxl_wrong_lora.safetensors',
              'SDXL_FILM_PHOTOGRAPHY_STYLE_BetaV0.4_139125.safetensors',
              'RMSDXL_Darkness_Cinema_211126.safetensors',
              'ASH_230253.safetensors',
              'xl_more_art-full_v1_113467.safetensors',
              'add-detail-xl_99264.safetensors',
              'Kodak Motion Picture Film Style3_222549.safetensors',
              'JuggerCineXL2_95870.safetensors',
              'AST_229222.safetensors',
              'NsfwPovAllInOneLoraSdxl-000009_120561.safetensors',
              'LCM_LoRA_Weights_169352.safetensors',
              'Cinematic_Painting_XL_V02_151050.safetensors',
              'SDXLHighDetail_v5_107780.safetensors',
              'p1c4ss0_003-step00028000_140158.safetensors',
              'RMSDXL_Enhance_209720.safetensors',
              'DI_belle_delphine_sdxl_v1_93586.safetensors',
              'LogoRedmondV2-Logo-LogoRedmAF_135681.safetensors',
              'AnalogRedmondV2-Analog-AnalogRedmAF_135776.safetensors',
              'csal_scenery_XL_106371.safetensors',
              'Holography_140985.safetensors',
              'SDXL_cinematic_Sa_May_151691.safetensors',
              'BoutXOval_SDXL_dim32-000005_92430.safetensors',
              'chahua_137600.safetensors',
              'Macpaint-XL-000015_131908.safetensors',
              'fdec_140071.safetensors',
              'landscape-painting-sdxl_v2_111037.safetensors',
              'SDXL-Emoji-Lora-r4_120592.safetensors',
              'vncnt_144477.safetensors',
              'Digital_Madness_139496.safetensors',
              'Aether_Glitch_v1_LoRA_128844.safetensors',
              'flat_chested_with_top-v0.9_141431.safetensors',
              'fr4z3tt4_94108.safetensors',
              'cute girl_104277.safetensors',
              'SDXL_MSPaint_Portrait_157817.safetensors',
              'Storyboard_sketch_140144.safetensors',
              '3l3ctronics-step00003000_126441.safetensors',
              'sdxl_glass_136034.safetensors',
              'sd_xl_dpo_lora_v1_214128.safetensors',
              'Crayon_Drawing_159195.safetensors',
              'concept_pov_dt_xl2-000020_119643.safetensors',
              'SDXLGhostStyle_145763.safetensors',
              'polyhedron_all_sdxl-000004_110557.safetensors',
              'xl_yoshiaki_kawajiri_v1r64_126468.safetensors',
              'Lines_Clothes_XL_138401.safetensors',
              'sdxl_cute_social_comic-000002_107980.safetensors',
              'Graphic_Portrait_146112.safetensors',
              'FF-Halloween-Full_154412.safetensors',
              'Candy_land_2_141469.safetensors',
              'acidzlime-sdxl_154149.safetensors',
              'NsfwPovAllInOneLoraSdxl-000009MINI_120545.safetensors',
              'sdxl_liminal_v1-step00003000_105309.safetensors',
              'Cute Animals_114772.safetensors',
              'ral-beer-sdxl_235173.safetensors',
              'JoeyKingSDXL_201358.safetensors',
              'SDXLFrosted_191661.safetensors',
              'gardevoir-000017_101442.safetensors',
              'Graphic_140518.safetensors',
              'TShirtDesignRedmondV2-Tshirtdesign-TshirtDesignAF_136145.safetensors',
              'Film Stock Footage Style_235607.safetensors',
              'ral-wtchz-sdxl_233487.safetensors',
              'Pussy_pants_nohands_D03_realvizXL2_bucket_01-000010_200586.safetensors',
              'cinematic vintage film_165734.safetensors',
              'epoxy_skull-sdxl_153213.safetensors']


def generate_image(novita_key, upload_images, checkpoint, lora_name, lora_strength, width, height, prompt, style_name, negative_prompt, num_steps, style_strength_ratio, num_outputs, guidance_scale, seed):
    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)

    if upload_images is None:
        raise gr.Error(f"Cannot find any input face image! Please refer to step 1️⃣")

    loras = []
    if lora_name:
        loras.append(
            MakePhotoLoRA(
                model_name=lora_name,
                strength=lora_strength
            )
        )

    try:
        res = get_noviata_client(novita_key).make_photo(
            model_name=checkpoint,
            prompt=prompt,
            negative_prompt=negative_prompt,
            images=[_.name for _ in upload_images],
            loras=loras,
            steps=num_steps,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            strength=style_strength_ratio / 100,
            seed=seed,
            image_num=num_outputs,
        )
    except Exception as e:
        raise gr.Error(f"Error: {e}")

    images = [Image.open(BytesIO(base64.b64decode(res.images_encoded[idx]))) for idx in range(len(res.images_encoded))]

    return images, {
        "model_name": checkpoint,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "images": [_.name for _ in upload_images],
        "loras": loras,
        "steps": num_steps,
        "height": height,
        "width": width,
        "guidance_scale": guidance_scale,
        "strength": style_strength_ratio / 100,
        "seed": seed,
        "image_num": num_outputs,
    }


def swap_to_gallery(images):
    return gr.update(value=[_.name for _ in images], visible=True), gr.update(visible=True), gr.update(visible=False)


def upload_example_to_gallery(images, checkpoint, lora_name, lora_strength, width, height, prompt, style, negative_prompt):
    return gr.update(value=[_.name for _ in images], visible=True), gr.update(value=checkpoint), gr.update(visible=True), gr.update(visible=False)


def remove_back_to_files():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)


def remove_tips():
    return gr.update(visible=False)


def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed


def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + ' ' + negative


def get_image_path_list(folder_name):
    image_basename_list = os.listdir(folder_name)
    image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])
    return image_path_list


def get_example():
    case = [
        [
            get_image_path_list('./examples/scarletthead_woman'),
            "protovisionXLHighFidelity3D_releaseV660Bakedvae_207131.safetensors",
            None,
            None,
            1024,
            1024,
            "portrait photo of a woman img, perfect face, natural skin,film grain,White dress,castle",
            "(No style)",
            "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",
        ],
        [
            get_image_path_list('./examples/newton_man'),
            "protovisionXLHighFidelity3D_releaseV660Bakedvae_207131.safetensors",
            None,
            None,
            1024,
            1024,
            "portrait photo of a man img, perfect face, natural skin,film grain,spider man,city,upper body ",
            "(No style)",
            "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",
        ],
    ]
    return case


with gr.Blocks() as demo:
    # gr.DuplicateButton(
    #     value="Duplicate Space for private use ",
    #     elem_id="duplicate-button",
    #     visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    # )
    gr.HTML(''' 
        <a href="https://novita.ai/?utm_source=huggingface&utm_medium=make-photo&utm_campaign=make-photo">
            <img src="https://raw.githubusercontent.com/wiki/novitalabs/sd-webui-cleaner/images/logo2.png" width="120px;" alt="Unsplash" />
        </a>
        <h1>Make Photo Playground</h1>
        <h3>Start integrate with <a href="https://novita.ai/get-started/UseCase_ImageEnhancement.html#_15-upscale?utm_source=huggingface&utm_medium=make-photo&utm_campaign=make-photo">Make Photo API</a> 
        <h3> Get Novita.AI API Key from <a href="https://novita.ai/get-started/Account_account_and_key.html?utm_source=huggingface&utm_medium=make-photo&utm_campaign=make-photo">Novita.AI</a></h2>
        '''
            )

    with gr.Row():
        with gr.Column(scale=1):
            novita_key = gr.Textbox(value="", label="Novita.AI API KEY (store in broweser)", placeholder="novita.ai api key", type="password")
        with gr.Column(scale=1):
            user_balance = gr.Textbox(label="User Balance", value="0.0")

    with gr.Row():
        with gr.Column():
            files = gr.Files(
                label="Drag (Select) 1 or more photos of your face",
                file_types=["image"]
            )
            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
            with gr.Column(visible=False) as clear_button:
                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
            checkpoint = gr.Dropdown(label="Checkpoint", choices=sdxl_models, value="protovisionXLHighFidelity3D_releaseV660Bakedvae_207131.safetensors", type="value")
            lora_name = gr.Dropdown(label="LoRA", choices=sdxl_loras)
            lora_strength = gr.Slider(label="LoRA strength", minimum=0.0, maximum=1.0, step=0.1, value=0.5)
            height = gr.Slider(label="Height", minimum=256, maximum=2048, step=1, value=1024)
            width = gr.Slider(label="Width", minimum=256, maximum=2048, step=1, value=1024)

            prompt = gr.Textbox(label="Prompt",
                                info="Try something like 'a photo of a man/woman img', 'img' is the trigger word.",
                                placeholder="A photo of a [man/woman img]...")
            style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
            submit = gr.Button("Submit")

            with gr.Accordion(open=False, label="Advanced Options"):
                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    placeholder="low quality",
                    value="nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry",
                )
                num_steps = gr.Slider(
                    label="Number of sample steps",
                    minimum=20,
                    maximum=100,
                    step=1,
                    value=50,
                )
                style_strength_ratio = gr.Slider(
                    label="Style strength (%)",
                    minimum=15,
                    maximum=50,
                    step=1,
                    value=20,
                )
                num_outputs = gr.Slider(
                    label="Number of output images",
                    minimum=1,
                    maximum=4,
                    step=1,
                    value=1,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Column():
            gallery = gr.Gallery(label="Generated Images")
            request_payload = gr.JSON(label="Request Payload")

        files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])

        submit.click(
            fn=remove_tips,
        ).then(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
        ).then(
            fn=generate_image,
            inputs=[novita_key, files, checkpoint, lora_name, lora_strength, width, height, prompt,  style,  negative_prompt, num_steps, style_strength_ratio, num_outputs, guidance_scale, seed],
            outputs=[gallery, request_payload]
        )

    gr.Examples(
        examples=get_example(),
        inputs=[files, checkpoint, lora_name, lora_strength, width, height, prompt, style, negative_prompt],
        run_on_click=True,
        fn=upload_example_to_gallery,
        outputs=[uploaded_files, checkpoint, clear_button, files],
    )

    def onload(novita_key):
        if novita_key is None or novita_key == "":
            return novita_key, f"$ UNKNOWN", gr.update(visible=False)
        try:
            user_info_json = get_noviata_client(novita_key).user_info()
        except Exception as e:
            return novita_key, f"$ UNKNOWN"

        return novita_key, f"$ {user_info_json.credit_balance / 100 / 100:.2f}"

    novita_key.change(onload, inputs=novita_key, outputs=[novita_key, user_balance], _js="(v)=>{ setStorage('novita_key',v); return [v]; }")

    demo.load(
        inputs=[novita_key],
        outputs=[novita_key, user_balance],
        fn=onload,
        _js=get_local_storage,
    )

demo.queue(api_open=False, concurrency_count=20)
demo.launch(server_name="0.0.0.0", share=False)