css and speed

app.py

@@ -12,6 +12,7 @@ import gradio as gr
 import numpy as np
 import spaces
 # import imageio
+import gc
 import torch
 from PIL import Image
 from diffusers import (
@@ -113,16 +114,10 @@ preprocessor = Preprocessor()
 preprocessor.load("NormalBae")
 
 print("---------------Loaded controlnet pipeline---------------") 
+torch.cuda.empty_cache()
+gc.collect()
 print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
-
-
-# @spaces.GPU(duration=12)
-# def init(pipe):
-#     pipe.enable_xformers_memory_efficient_attention()
-#     pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-#     pipe.unet.set_attn_processor(AttnProcessor2_0())
-#     print("Model Compiled!")
-# init(pipe)
+print("Model Compiled!")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -246,20 +241,30 @@ def apply_style(style_name):
 
     
 css = """
-h1 {
+h1, h2, h3 {
     text-align: center;
-    display:block;
+    display: block;
 }
-h2 {
-    text-align: center;
-    display:block;
+footer {
+    visibility: hidden;
 }
-h3 {
-    text-align: center;
-    display:block;
+.gradio-container {
+    max-width: 900px !important;
+}
+.gr-image {
+    display: flex;
+    justify-content: center; 
+    align-items: center;
+    width: 100%;
+    height: 512px;
+    overflow: hidden;
+}
+.gr-image img {
+    width: 100%;
+    height: 100%; 
+    object-fit: cover;
+    object-position: center;
 }
-.gradio-container{max-width: 1200px !important}
-footer {visibility: hidden}
 """
 with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     #############################################################################
@@ -400,7 +405,7 @@ def process_image(
     guidance_scale,
     seed,
 ):
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     preprocess_start = time.time()
     print("processing image")
 
@@ -439,8 +444,8 @@ def process_image(
     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()
+    # torch.cuda.synchronize()
+    # torch.cuda.empty_cache()
     return results
 
 if prod:

local_app.py

@@ -10,7 +10,7 @@ import random
 import time
 import gradio as gr
 import numpy as np
-# import imageio
+import gc
 import torch
 from PIL import Image
 from diffusers import (
@@ -19,7 +19,7 @@ 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)
 
@@ -58,9 +58,9 @@ if gr.NO_RELOAD:
     base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
     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 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(
@@ -69,7 +69,7 @@ if gr.NO_RELOAD:
         # load_safety_checker=True,
         controlnet=controlnet,
         scheduler=scheduler,
-        vae=vae,
+        # vae=vae,
         torch_dtype=torch.float16,
     )
 
@@ -88,10 +88,14 @@ if gr.NO_RELOAD:
     print("loading preprocessor")
     from preprocess import Preprocessor
     preprocessor = Preprocessor()
-    # preprocessor.load("NormalBae")
+    preprocessor.load("NormalBae")
 
     print("---------------Loaded controlnet pipeline---------------") 
+    pipe.unet.set_attn_processor(AttnProcessor2_0())
+    torch.cuda.empty_cache()
+    gc.collect()
     print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
+    print("Model Compiled!")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -219,12 +223,12 @@ h1, h2, h3 {
     text-align: center;
     display: block;
 }
-.gradio-container {
-    max-width: 1200px !important;
-}
 footer {
     visibility: hidden;
 }
+.gradio-container {
+    max-width: 900px !important;
+}
 .gr-image {
     display: flex;
     justify-content: center; 
@@ -293,8 +297,8 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             label="Design Styles",
         )
     # input image
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=1, min_width=300):
+    with gr.Row():
+        with gr.Column():
             image = gr.Image(
                 label="Input",
                 sources=["upload"],
@@ -306,7 +310,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             with gr.Column():
                 run_button = gr.Button(value="Use this one", size="lg", visible=False)
         # output image
-        with gr.Column(scale=1, min_width=300):
+        with gr.Column():
             result = gr.Image(  
                 label="Output",
                 interactive=False,
@@ -338,18 +342,21 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     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 and change input
-    @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=[image, result], show_progress="minimal")
+    # 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")
-            image = Image.open("temp_image.jpg")
-            yield gr.update(value=image), gr.update()
+            ai_temp_image = Image.open("temp_image.jpg")
+            return ai_temp_image
         except FileNotFoundError:
             print("No AI Image Available")
-            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
+            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")
@@ -375,7 +382,7 @@ def process_image(
     guidance_scale,
     seed,
 ):
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     preprocess_start = time.time()
     print("processing image")
     preprocessor.load("NormalBae")
@@ -400,21 +407,20 @@ def process_image(
     negative_prompt=str(n_prompt)
     print(prompt)
     start = time.time()
-    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]
+    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()
+    # torch.cuda.synchronize()
+    # torch.cuda.empty_cache()
     return results
 
 if prod: