Update app.py

app.py

@@ -10,6 +10,7 @@ import urllib
 import time
 import os
 import datetime
+import gc
 
 from models.transformer_sd3 import SD3Transformer2DModel
 #from diffusers import StableDiffusion3Pipeline
@@ -34,7 +35,6 @@ torch.backends.cudnn.deterministic = False
 torch.backends.cudnn.benchmark = False
 #torch.backends.cuda.preferred_blas_library="cublas"
 #torch.backends.cuda.preferred_linalg_library="cusolver"
-torch.set_float32_matmul_precision("highest")
 
 hftoken = os.getenv("HF_TOKEN") 
 
@@ -57,42 +57,37 @@ def upload_to_ftp(filename):
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 torch_dtype = torch.bfloat16
 
-def load_and_prepare_models():
-    transformer = SD3Transformer2DModel.from_pretrained(
-        model_path, subfolder="transformer" #, torch_dtype=torch.bfloat16
-    )
-    vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", safety_checker=None, use_safetensors=True, low_cpu_mem_usage=False, subfolder='vae', torch_dtype=torch.float32, token=True)
-    pipe = StableDiffusion3Pipeline.from_pretrained(
-        #"stabilityai  #  stable-diffusion-3.5-large",
-        "ford442/stable-diffusion-3.5-large-bf16",
-         #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
-        text_encoder=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
-        text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
-        text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
-        #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
-        #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
-        tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", use_fast=True, subfolder="tokenizer_3", token=True),
-        #torch_dtype=torch.bfloat16,
-        transformer=transformer,
-        vae=None
-        #use_safetensors=False,
-    )
-    torch.cuda.empty_cache()
-    pipe.to(device=device, dtype=torch.bfloat16)
-    pipe.vae=vaeX.to(device)
-    upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
-    torch.cuda.empty_cache()
-    return pipe, upscaler
+transformer = SD3Transformer2DModel.from_pretrained(
+    model_path, subfolder="transformer" #, torch_dtype=torch.bfloat16
+)
 
-text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
-text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
-text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
+vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", safety_checker=None, use_safetensors=True, low_cpu_mem_usage=False, subfolder='vae', torch_dtype=torch.float32, token=True)
+
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    #"stabilityai  #  stable-diffusion-3.5-large",
+    "ford442/stable-diffusion-3.5-large-bf16",
+     #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
+    text_encoder=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
+    text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
+    text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
+    #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
+    #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
+    tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", use_fast=True, subfolder="tokenizer_3", token=True),
+    #torch_dtype=torch.bfloat16,
+    transformer=transformer,
+    vae=None
+    #use_safetensors=False,
+)
 
-pipe, upscaler_2 = load_and_prepare_models()
+pipe.to(device=device, dtype=torch.bfloat16)
 
 #pipe.to(device)
+pipe.vae=vaeX.to(device)
+text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
+text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
+text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
 
-#upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
+upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
@@ -131,8 +126,10 @@ def infer(
         nb_token=64, 
     )
     upscaler_2.to(torch.device('cpu'))
+    gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.reset_peak_memory_stats()
+    torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
     enhanced_prompt = prompt
@@ -190,6 +187,9 @@ def infer(
         sd_image.save(rv_path,optimize=False,compress_level=0)
         upload_to_ftp(rv_path)
         upscaler_2.to(torch.device('cuda'))
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
         with torch.no_grad():
             upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
         print('-- got upscaled image --')
@@ -219,7 +219,7 @@ body{
 
 with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # StableDiffusion 3.5 Large with IP Adapter Test B")
+        gr.Markdown(" # StableDiffusion 3.5 Large with IP Adapter")
         expanded_prompt_output = gr.Textbox(label="Prompt", lines=5) 
         with gr.Row():
             prompt = gr.Text(
@@ -281,7 +281,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                     value=1.0,
                 )
                 image_encoder_path = gr.Dropdown(
-                    ["google/siglip-so400m-patch14-384", "google/siglip-base-patch16-512", "jancuhel/google-siglip-so400m-patch14-384-img-text-relevancy", "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"], 
+                    ["google/siglip-so400m-patch14-384", "jancuhel/google-siglip-so400m-patch14-384-img-text-relevancy", "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"], 
                     label="CLIP Model", 
                 )
                 ip_scale = gr.Slider(