cleanup

.gitignore

@@ -5,4 +5,5 @@ anime_app_local.py
 *.jpg
 *.png
 *.pyc
-client.py
+client.py
+outputs/*

Dockerfile

@@ -33,8 +33,8 @@ ENV HOME=/home/user \
 
 RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
 
-RUN apt-get install -y nvidia-340
-RUN nvidia-smi
+#RUN apt-get install -y nvidia-340
+#RUN nvidia-smi
 
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app

anime_app.py

@@ -1,6 +1,8 @@
 prod = True
+port = 8080
 show_options = True
 if prod:
+    port = 8081
     show_options = False
 import os
 import gc
@@ -8,6 +10,7 @@ import random
 import time
 import gradio as gr
 import spaces
+import numpy as np
 import imageio
 from huggingface_hub import HfApi
 import torch
@@ -18,8 +21,7 @@ from diffusers import (
     StableDiffusionControlNetPipeline,
 )
 from preprocess_anime import Preprocessor
-# MAX_SEED = np.iinfo(np.int32).max
-MAX_SEED = 2147483647
+MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
 print("CUDA version:", torch.version.cuda)
@@ -98,6 +100,7 @@ def get_additional_prompt():
 
 def get_prompt(prompt, additional_prompt):
     default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     randomize = get_additional_prompt()
     # nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     # bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
@@ -113,17 +116,18 @@ def get_prompt(prompt, additional_prompt):
     naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
     abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
     # shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari"
-    shibari2 = "extremely detailed, hyperrealistic photography,earrings,tattoo,score_9, HDA_Shibari"
+    shibari2 = "octane render, highly detailed, volumetric, HDA_Shibari"
     
     if prompt == "":
-        prompts = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
+        prompts = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2]
         prompts_nsfw = [abg, shibari2, ahegao2]
         preset = random.choice(prompts)
         prompt = f"{preset}"
         # print(f"-------------{preset}-------------")
     else:
         # prompt = f"{prompt}, {randomize}"
-        prompt = f"{default},{prompt}"
+        # prompt = f"{default},{prompt}"
+        prompt = default2
     print(f"{prompt}")
     return prompt
 
@@ -237,7 +241,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 step=1,
             )
             num_steps = gr.Slider(
-                label="Number of steps", minimum=1, maximum=100, value=12, step=1
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
             )  # 20/4.5 or 12 without lora, 4 with lora
             guidance_scale = gr.Slider(
                 label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
@@ -267,7 +271,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 sources=["upload"],
                 show_label=True,
                 mirror_webcam=True,
-                format="jpg",
+                format="webp",
             )
             # run button
             with gr.Column():
@@ -277,7 +281,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             result = gr.Image(  
                 label="Anime AI",
                 interactive=False,
-                format="jpg",
+                format="webp",
                 show_share_button= False,
             )
             # Use this image button
@@ -330,4 +334,5 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     def turn_buttons_on():
         return gr.update(visible=True), gr.update(visible=True)
 
-demo.queue(api_open=False).launch(show_api=False)
+# demo.queue(api_open=False).launch(show_api=False)
+demo.queue(max_size=20).launch(server_name="localhost", server_port=port)

profiler.py

@@ -1,381 +0,0 @@
-import cProfile
-import pstats
-import io
-import gc
-import random
-import time
-import gradio as gr
-import spaces
-import imageio
-from huggingface_hub import HfApi
-import torch
-from PIL import Image
-from diffusers import (
-    ControlNetModel,
-    DPMSolverMultistepScheduler,
-    StableDiffusionControlNetPipeline,
-)
-from preprocess_anime import Preprocessor, conditionally_manage_memory
-from settings import API_KEY, MAX_NUM_IMAGES, MAX_SEED
-
-preprocessor = None
-controlnet = None
-scheduler = None
-pipe = None
-compiled = False
-api = HfApi()
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-def get_additional_prompt():
-    prompt = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
-    top = ["tank top", "blouse", "button up shirt", "sweater", "corset top"]
-    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt", "leggings", "high-waisted shorts"]
-    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker", "necklace", "headband", "headphones"]
-    return f"{prompt}, {random.choice(top)}, {random.choice(bottom)}, {random.choice(accessory)}, score_9"
-    # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]
-
-def get_prompt(prompt, additional_prompt):
-    default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
-    randomize = get_additional_prompt()
-    nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
-    bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
-    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
-    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
-    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
-    ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
-    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
-    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
-    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
-    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
-    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
-    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
-    abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
-    shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari"
-        
-    if prompt == "":
-        prompts = [randomize, nude, bodypaint, pet_play, bondage, ahegao2, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari]
-        prompts_nsfw = [nude, bodypaint, abg, ahegao2, shibari]
-        preset = random.choice(prompts)
-        prompt = f"{preset}"
-        # print(f"-------------{preset}-------------")
-    else:
-        # prompt = f"{prompt}, {randomize}"
-        prompt = f"{default},{prompt}"
-    print(f"{prompt}")
-    return prompt
-
-def initialize_models():
-    global preprocessor, controlnet, scheduler, pipe
-    if preprocessor is None:
-        preprocessor = Preprocessor()
-
-    if controlnet is None:
-        model_id = "lllyasviel/control_v11p_sd15_normalbae"
-        print("initializing controlnet")
-        controlnet = ControlNetModel.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2",
-        ).to("cuda")
-
-    if scheduler is None:
-        scheduler = DPMSolverMultistepScheduler.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
-            solver_order=2,
-            subfolder="scheduler",
-            use_karras_sigmas=True,
-            final_sigmas_type="sigma_min",
-            algorithm_type="sde-dpmsolver++",
-            prediction_type="epsilon",
-            thresholding=False,
-            denoise_final=True,
-            device_map="cuda",
-        )
-
-    if pipe is None:
-        base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
-        pipe = StableDiffusionControlNetPipeline.from_single_file(
-            base_model_url,
-            safety_checker=None,
-            controlnet=controlnet,
-            scheduler=scheduler,
-            torch_dtype=torch.float16,
-        )
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
-        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
-        pipe.to("cuda")
-        print("---------------Loaded controlnet pipeline---------------")
-
-@spaces.GPU(duration=11)
-@torch.inference_mode()
-def process_image(
-    image,
-    prompt,
-    a_prompt,
-    n_prompt,
-    num_images,
-    image_resolution,
-    preprocess_resolution,
-    num_steps,
-    guidance_scale,
-    seed,
-):
-    initialize_models()
-    preprocessor.load("NormalBae")
-    control_image = preprocessor(
-        image=image,
-        image_resolution=image_resolution,
-        detect_resolution=preprocess_resolution,
-    )
-    custom_prompt = str(get_prompt(prompt, a_prompt))
-    negative_prompt = str(n_prompt)
-    global compiled
-    generator = torch.cuda.manual_seed(seed)
-    if not compiled:
-        print("-----------------------------------Not Compiled-----------------------------------")
-        compiled = True
-    start = time.time()
-    results = pipe(
-        prompt=custom_prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=num_images,
-        num_inference_steps=num_steps,
-        generator=generator,
-        image=control_image,
-    ).images[0]
-    print(f"Inference done in: {time.time() - start:.2f} seconds")
-    
-    timestamp = int(time.time())
-    img_path = f"{timestamp}.jpg"
-    results_path = f"{timestamp}_out.jpg"
-    imageio.imsave(img_path, image)
-    results.save(results_path)
-    
-    api.upload_file(
-        path_or_fileobj=img_path,
-        path_in_repo=img_path,
-        repo_id="broyang/anime-ai-outputs",
-        repo_type="dataset",
-        token=API_KEY,
-        run_as_future=True,
-    )
-    api.upload_file(
-        path_or_fileobj=results_path,
-        path_in_repo=results_path,
-        repo_id="broyang/anime-ai-outputs",
-        repo_type="dataset",
-        token=API_KEY,
-        run_as_future=True,
-    )
-
-    conditionally_manage_memory()
-
-    results.save("temp_image.png")
-    return results
-
-def main():
-    prod = True
-    show_options = True
-    if prod:
-        show_options = False
-
-    print("CUDA version:", torch.version.cuda)
-    print("loading pipe")
-
-    css = """
-    h1 {
-        text-align: center;
-        display:block;
-    }
-    h2 {
-        text-align: center;
-        display:block;
-    }
-    h3 {
-        text-align: center;
-        display:block;
-    }
-    footer {visibility: hidden}
-    """
-    with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-        with gr.Row():
-            with gr.Accordion("Advanced options", open=show_options, visible=show_options):
-                num_images = gr.Slider(
-                    label="Images", minimum=1, maximum=MAX_NUM_IMAGES, value=1, step=1
-                )
-                image_resolution = gr.Slider(
-                    label="Image resolution",
-                    minimum=256,
-                    maximum=1024,
-                    value=768,
-                    step=256,
-                )
-                preprocess_resolution = gr.Slider(
-                    label="Preprocess resolution",
-                    minimum=128,
-                    maximum=1024,
-                    value=768,
-                    step=1,
-                )
-                num_steps = gr.Slider(
-                    label="Number of steps", minimum=1, maximum=100, value=12, step=1
-                )
-                guidance_scale = gr.Slider(
-                    label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
-                )
-                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                a_prompt = gr.Textbox(
-                    label="Additional prompt",
-                    value = ""
-                )
-                n_prompt = gr.Textbox(
-                    label="Negative prompt",
-                    value="EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
-                )
-        with gr.Column():
-            prompt = gr.Textbox(
-                label="Description",
-                placeholder="Leave empty for something spicy 👀",
-            )
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image(
-                    label="Input",
-                    sources=["upload"],
-                    show_label=True,
-                    format="webp",
-                )
-                with gr.Column():
-                    run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
-            with gr.Column():
-                result = gr.Image(
-                    label="Anime AI",
-                    interactive=False,
-                    format="webp",
-                    visible = True,
-                    show_share_button= False,
-                )
-                with gr.Column():
-                    use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
-        config = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_images,
-            image_resolution,
-            preprocess_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-        ]
-        
-        @spaces.GPU(duration=11)
-        @torch.inference_mode()
-        @gr.on(triggers=[image.upload], inputs=config, outputs=[result])
-        def auto_process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed):
-            return process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-    
-        @gr.on(triggers=[image.upload], inputs=None, outputs=[use_ai_button, run_button])
-        def turn_buttons_off():
-            return gr.update(visible=False), gr.update(visible=False)
-    
-        @gr.on(triggers=[use_ai_button.click], inputs=None, outputs=[use_ai_button, run_button])
-        def turn_buttons_off():
-            return gr.update(visible=False), gr.update(visible=False)
-    
-        @gr.on(triggers=[run_button.click], inputs=None, outputs=[use_ai_button, run_button])
-        def turn_buttons_off():
-            return gr.update(visible=False), gr.update(visible=False)
-    
-        @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button])
-        def turn_buttons_on():
-            return gr.update(visible=True), gr.update(visible=True)
-        
-        with gr.Row():
-            helper_text = gr.Markdown("## Tap and hold (on mobile) to save the image.", visible=True)
-            
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-            show_progress="none",
-        ).then(
-            fn=auto_process_image,
-            inputs=config,
-            outputs=result,
-            api_name=False,
-            show_progress="minimal",
-        )
-
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-            show_progress="none",
-        ).then(
-            fn=auto_process_image,
-            inputs=config,
-            outputs=result,
-            show_progress="minimal",
-        )
-
-        def update_config():
-            try:
-                print("Updating image to AI Temp Image")
-                ai_temp_image = Image.open("temp_image.png")
-                return ai_temp_image
-            except FileNotFoundError:
-                print("No AI Image Available")
-                return None
-        
-        use_ai_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-            show_progress="none",
-        ).then(
-            fn=lambda _: update_config(),
-            inputs=[image],
-            outputs=image,
-            show_progress="minimal",
-        ).then(
-            fn=auto_process_image,
-            inputs=[image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed],
-            outputs=result,
-            show_progress="minimal",
-        )
-
-    demo.launch()
-
-if __name__ == "__main__":
-    pr = cProfile.Profile()
-    pr.enable()
-    main()
-    pr.disable()
-
-    s = io.StringIO()
-    sortby = 'cumulative'
-    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
-    ps.print_stats()
-    print(s.getvalue())