refactor

anime_app.py

@@ -2,6 +2,7 @@ prod = True
 show_options = True
 if prod:
     show_options = False
+import os
 import gc
 import random
 import time
@@ -17,7 +18,9 @@ from diffusers import (
     StableDiffusionControlNetPipeline,
 )
 from preprocess_anime import Preprocessor
-from settings import API_KEY, MAX_NUM_IMAGES, MAX_SEED
+# MAX_SEED = np.iinfo(np.int32).max
+MAX_SEED = 2147483647
+API_KEY = os.environ.get("API_KEY", None)
 
 print("CUDA version:", torch.version.cuda)
 print("loading pipe")
@@ -102,7 +105,7 @@ def get_prompt(prompt, additional_prompt):
     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"
+    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"
@@ -113,8 +116,8 @@ def get_prompt(prompt, additional_prompt):
     shibari2 = "extremely detailed, hyperrealistic photography,earrings,tattoo,score_9, HDA_Shibari"
     
     if prompt == "":
-        prompts = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2]
-        prompts_nsfw = [abg, shibari2]
+        prompts = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
+        prompts_nsfw = [abg, shibari2, ahegao2]
         preset = random.choice(prompts)
         prompt = f"{preset}"
         # print(f"-------------{preset}-------------")
@@ -124,6 +127,7 @@ def get_prompt(prompt, additional_prompt):
     print(f"{prompt}")
     return prompt
 
+
 @spaces.GPU(duration=20)
 @torch.inference_mode()
 def process_image(
@@ -137,10 +141,10 @@ def process_image(
     num_steps,
     guidance_scale,
     seed,
+    progress=gr.Progress(track_tqdm=True)
 ):
     print("processing image")
     start = time.time()
-    
     preprocessor.load("NormalBae")
     # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
     
@@ -152,6 +156,7 @@ def process_image(
     custom_prompt=str(get_prompt(prompt, a_prompt))
     negative_prompt=str(n_prompt)
     global compiled
+    seed = random.randint(0, MAX_SEED)  
     generator = torch.cuda.manual_seed(seed)
     if not compiled:
         print("Not Compiled")
@@ -168,8 +173,8 @@ def process_image(
     print(f"\n-------------------------Processed in: {time.time() - start:.2f} seconds-------------------------")
     
     timestamp = int(time.time())
-    img_path = f"{timestamp}.jpg"
-    results_path = f"{timestamp}_out_{prompt}.jpg"
+    img_path = f"./outputs/{timestamp}.jpg"
+    results_path = f"./outputs/{timestamp}_out_{prompt}.jpg"
     imageio.imsave(img_path, image)
     results.save(results_path)
     results.save("temp_image.jpg")
@@ -190,9 +195,6 @@ def process_image(
         token=API_KEY,
         run_as_future=True,
     )
-    
-    # torch.cuda.empty_cache()
-    # gc.collect()
 
     return results
 
@@ -216,7 +218,7 @@ 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
+                label="Images", minimum=1, maximum=4, value=1, step=1
             )
             image_resolution = gr.Slider(
                 label="Image resolution",
@@ -274,7 +276,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 label="Anime AI",
                 interactive=False,
                 format="webp",
-                visible = True,
                 show_share_button= False,
             )
             # Use this image button
@@ -293,89 +294,38 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         seed,
     ]
     
-    @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):
-        seed = randomize_seed_fn(seed, randomize_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=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=process_image,
-        inputs=config,
-        outputs=result,
-        show_progress="minimal",
-    )
+    
+    # image processing
+    @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config, outputs=result, show_progress="minimal")
+    def auto_process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+        return process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
+    
+    # AI Image Processing
+    @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=result, show_progress="minimal")
+    def submit(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+        return process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
 
-    def update_config():
+    # Change input to result
+    @gr.on(triggers=[use_ai_button.click], inputs=None, outputs=image, show_progress="hidden")
+    def update_input():
         try:
             print("Updating image to AI Temp Image")
-            # Read the image from the file
             ai_temp_image = Image.open("temp_image.jpg")
             return ai_temp_image
         except FileNotFoundError:
             print("No AI Image Available")
             return None
+        
+    # Turn off buttons when processing
+    @gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None, outputs=[run_button, use_ai_button], show_progress="hidden")
+    def turn_buttons_off():
+        return gr.update(visible=False), gr.update(visible=False)
     
-    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,
-        queue=False,
-        show_progress="minimal",
-    ).then(
-        fn=process_image,
-        inputs=[image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed],
-        outputs=result,
-        queue=False,
-        show_progress="minimal",
-    )
+    # Turn on buttons when processing is complete
+    @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button], show_progress="hidden")
+    def turn_buttons_on():
+        return gr.update(visible=True), gr.update(visible=True)
 
-    
 demo.queue(api_open=False).launch(show_api=False)

cv_utils.py

@@ -1,17 +0,0 @@
-import cv2
-import numpy as np
-
-
-def resize_image(input_image, resolution, interpolation=None):
-    H, W, C = input_image.shape
-    H = float(H)
-    W = float(W)
-    k = float(resolution) / max(H, W)
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    if interpolation is None:
-        interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
-    img = cv2.resize(input_image, (W, H), interpolation=interpolation)
-    return img

preprocess_anime.py

@@ -6,10 +6,26 @@ import torch
 from controlnet_aux import NormalBaeDetector#, CannyDetector
 
 from controlnet_aux.util import HWC3
-from cv_utils import resize_image
+import cv2
+# from cv_utils import resize_image
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
+    
+    def resize_image(input_image, resolution, interpolation=None):
+        H, W, C = input_image.shape
+        H = float(H)
+        W = float(W)
+        k = float(resolution) / max(H, W)
+        H *= k
+        W *= k
+        H = int(np.round(H / 64.0)) * 64
+        W = int(np.round(W / 64.0)) * 64
+        if interpolation is None:
+            interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
+        img = cv2.resize(input_image, (W, H), interpolation=interpolation)
+        return img
+
 
     def __init__(self):
         self.model = None

settings.py

@@ -1,16 +0,0 @@
-import os
-
-import numpy as np
-
-DEFAULT_MODEL_ID = os.getenv("DEFAULT_MODEL_ID", "runwayml/stable-diffusion-v1-5")
-# DEFAULT_MODEL_ID = ("runwayml/stable-diffusion-v1-5")
-MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "4"))
-DEFAULT_NUM_IMAGES = min(MAX_NUM_IMAGES, int(os.getenv("DEFAULT_NUM_IMAGES", "1")))
-MAX_IMAGE_RESOLUTION = int(os.getenv("MAX_IMAGE_RESOLUTION", "1024"))
-DEFAULT_IMAGE_RESOLUTION = min(MAX_IMAGE_RESOLUTION, int(os.getenv("DEFAULT_IMAGE_RESOLUTION", "512")))
-
-ALLOW_CHANGING_BASE_MODEL = os.getenv("SPACE_ID") != "hysts/ControlNet-v1-1"
-SHOW_DUPLICATE_BUTTON = os.getenv("SHOW_DUPLICATE_BUTTON") == "1"
-
-MAX_SEED = np.iinfo(np.int32).max
-API_KEY = os.environ["API_KEY"]

utils.py

@@ -1,9 +0,0 @@
-import random
-
-from settings import MAX_SEED
-
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed