local

app.py

@@ -351,28 +351,6 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             yield None, gr.update()
         result = process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
         yield gr.update(), result
-            
-    # # 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")
-    #         ai_temp_image = Image.open("temp_image.jpg")
-    #         return ai_temp_image
-    #     except FileNotFoundError:
-    #         print("No AI Image Available")
-    #         return None
-        
-    # # AI Image Processing
-    # @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=[result, image], show_progress="minimal")
-    # def submit(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
-    #     try:
-    #         print("Updating image to AI Temp Image")
-    #         image = Image.open("temp_image.jpg")
-    #     except FileNotFoundError:
-    #         print("No AI Image Available")
-    #         return None
-    #     return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed), image
 
     # 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")

local_app.py

@@ -6,7 +6,6 @@ if prod:
     # show_options = False
 
 import os
-import gc
 import random
 import time
 import gradio as gr
@@ -20,20 +19,16 @@ from diffusers import (
     StableDiffusionControlNetPipeline,
     AutoencoderKL,
 )
-from diffusers.models.attention_processor import AttnProcessor2_0
+# from diffusers.models.attention_processor import AttnProcessor2_0
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
 print("CUDA version:", torch.version.cuda)
-print("loading pipe")
+print("loading everything")
 compiled = False
 
 if gr.NO_RELOAD:
-    from preprocess import Preprocessor
-    preprocessor = Preprocessor()
-    # preprocessor.load("NormalBae")
     torch.cuda.max_memory_allocated(device="cuda")
-
     # Controlnet Normal
     model_id = "lllyasviel/control_v11p_sd15_normalbae"
     print("initializing controlnet")
@@ -42,37 +37,41 @@ if gr.NO_RELOAD:
         torch_dtype=torch.float16,
         attn_implementation="flash_attention_2",
     ).to("cuda")
-    
+
     # Scheduler
     scheduler = DPMSolverMultistepScheduler.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0", 
+        "runwayml/stable-diffusion-v1-5",
+        solver_order=2,
         subfolder="scheduler",
         use_karras_sigmas=True,
-        # final_sigmas_type="sigma_min",
+        final_sigmas_type="sigma_min",
         algorithm_type="sde-dpmsolver++",
-        # prediction_type="epsilon",
-        # thresholding=False,
+        prediction_type="epsilon",
+        thresholding=False,
         denoise_final=True,
         device_map="cuda",
-        attn_implementation="flash_attention_2",   
+        torch_dtype=torch.float16,
     )
 
     # Stable Diffusion Pipeline URL
     # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
     base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
-    base_model_id = "Lykon/absolute-reality-1.81"
     vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
-    
+
+    print('loading vae')
     vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
     vae.to(memory_format=torch.channels_last) 
 
+    print('loading pipe')
     pipe = StableDiffusionControlNetPipeline.from_single_file(
         base_model_url,
+        safety_checker=None,
+        # load_safety_checker=True,
         controlnet=controlnet,
         scheduler=scheduler,
         vae=vae,
         torch_dtype=torch.float16,
-    ).to("cuda")
+    )
 
     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")
@@ -85,15 +84,15 @@ if gr.NO_RELOAD:
     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("loading preprocessor")
+    from preprocess import Preprocessor
+    preprocessor = Preprocessor()
+    # preprocessor.load("NormalBae")
+
     print("---------------Loaded controlnet pipeline---------------") 
+    print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
 
-    # @spaces.GPU(duration=12)
-    pipe.unet.set_attn_processor(AttnProcessor2_0())
-    torch.cuda.empty_cache()
-    gc.collect()
-    print("Model Compiled!")
-    
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -132,7 +131,7 @@ def get_prompt(prompt, additional_prompt):
         girls = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
         prompts_nsfw = [abg, shibari2, ahegao2]
         prompt = f"{random.choice(girls)}"
-        # prompt = f"boho chic"
+        prompt = f"boho chic"
         # print(f"-------------{preset}-------------")
     else:
         prompt = f"Photo from Pinterest of {prompt} {interior}"
@@ -149,7 +148,7 @@ style_list = [
         "prompt": "Minimalist interior design, clean lines, neutral colors, uncluttered space, functional furniture, lots of natural light"
     },
     {
-        "name": "Boho Chic",
+        "name": "Boho",
         "prompt": "Bohemian chic interior, eclectic mix of patterns and textures, vintage furniture, plants, woven textiles, warm earthy colors"
     },
     {
@@ -213,40 +212,27 @@ def apply_style(style_name):
     if style_name in styles:
         p = styles.get(style_name, "boho chic")
     return p
-        
+
+    
 css = """
-h1, h2, h3 {
+h1 {
     text-align: center;
-    display: block;
-}
-
-.gradio-container {
-    max-width: 1200px !important;
+    display:block;
 }
-
-footer {
-    visibility: hidden;
-}
-
-.gr-image {
-    display: flex;
-    justify-content: center; 
-    align-items: center;
-    width: 100%;
-    height: 512px;
-    overflow: hidden;
+h2 {
+    text-align: center;
+    display:block;
 }
-
-.gr-image img {
-    width: 100%;
-    height: 100%; 
-    object-fit: cover;
-    object-position: center;
+h3 {
+    text-align: center;
+    display:block;
 }
+.gradio-container{max-width: 1200px !important}
+footer {visibility: hidden}
 """
 with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     #############################################################################
-    with gr.Row(): 
+    with gr.Row():
         with gr.Accordion("Advanced options", open=show_options, visible=show_options):
             num_images = gr.Slider(
                 label="Images", minimum=1, maximum=4, value=1, step=1
@@ -275,23 +261,20 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             a_prompt = gr.Textbox(
                 label="Additional prompt",
-                value = "design-style interior designed (interior space), captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length, tungsten white balance, (sharp focus), professional photography, high-resolution, 8k, Pulitzer Prize-winning"
+                value = "design-style interior designed (interior space), captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length, tungsten white balance"
             )
             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)",
             )
     #############################################################################
-    # input text
-    # with gr.Row():
-    #     gr.Text(label="Interior Design Style Examples", value="Eclectic, Maximalist, Bohemian, Scandinavian, Minimalist, Rustic, Modern Farmhouse, Contemporary, Luxury, Airbnb, Boho Chic, Midcentury Modern, Art Deco, Zen, Beach, Neoclassical, Industrial, Biophilic, Eco-friendly, Hollywood Glam, Parisian White, Saudi Prince Gold, French Country, Monster Energy Drink, Cyberpunk, Vaporwave, Baroque, etc.\n\nPro tip: add a color to customize it! You can also describe the furniture type.")
     with gr.Column():
         prompt = gr.Textbox(
-            label="Custom Prompt (optional)",
-            placeholder="use your imagination 👀",
+            label="Custom Prompt",
+            placeholder="boho chic",
         )
-    with gr.Column(visible=True):
-        style_selection = gr.Dropdown(
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
             show_label=True,
             container=True,
             interactive=True,
@@ -299,8 +282,9 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             value="None",
             label="Design Styles",
         )
-    with gr.Row(equal_height=True, variant="panel"): 
-        with gr.Column(min_width=300):
+    # input image
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1, min_width=300):
             image = gr.Image(
                 label="Input",
                 sources=["upload"],
@@ -308,17 +292,20 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 mirror_webcam=True,
                 format="webp",
             )
+            # run button
             with gr.Column():
-                run_button = gr.Button(value="Use this one", size="sm", visible=False)
-        with gr.Column(min_width=300):
+                run_button = gr.Button(value="Use this one", size="lg", visible=False)
+        # output image
+        with gr.Column(scale=1, min_width=300):
             result = gr.Image(  
                 label="Output",
                 interactive=False,
                 format="webp",
                 show_share_button= False,
             )
+            # Use this image button
             with gr.Column():
-                use_ai_button = gr.Button(value="Use this one", size="sm", visible=False)
+                use_ai_button = gr.Button(value="Use this one", size="lg", visible=False)
     config = [
         image,
         style_selection,
@@ -340,23 +327,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config, outputs=result, show_progress="minimal")
     def auto_process_image(image, style_selection, 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, style_selection, 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, style_selection, 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, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
 
-    # # 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")
-    #         ai_temp_image = Image.open("temp_image.jpg")
-    #         return ai_temp_image
-    #     except FileNotFoundError:
-    #         print("No AI Image Available")
-    #         return None
-        
     # AI Image Processing and change input
     @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=[image, result], show_progress="minimal")
     def submit(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
@@ -369,7 +340,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             yield None, gr.update()
         result = process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
         yield gr.update(), result
-        
+
     # 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():
@@ -380,7 +351,6 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     def turn_buttons_on():
         return gr.update(visible=True), gr.update(visible=True)
 
-# @spaces.GPU(duration=12)
 @torch.inference_mode()
 def process_image(
     image,
@@ -394,19 +364,17 @@ def process_image(
     num_steps,
     guidance_scale,
     seed,
-    progress=gr.Progress(track_tqdm=True)
 ):
     torch.cuda.synchronize()
     preprocess_start = time.time()
     print("processing image")
     preprocessor.load("NormalBae")
-    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
-    
+
     global compiled
     if not compiled:
         print("Not Compiled")
         compiled = True
-        
+    
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)
     control_image = preprocessor(
@@ -422,21 +390,21 @@ def process_image(
     negative_prompt=str(n_prompt)
     print(prompt)
     start = time.time()
-    results = pipe(
-        prompt=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]
+    with torch.cuda.amp.autocast():
+        results = pipe(
+            prompt=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"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
     results.save("temp_image.jpg")
     torch.cuda.synchronize()
     torch.cuda.empty_cache()
-    results.save("temp_image.jpg")
     return results
 
 if prod: