Update app.py

app.py

@@ -30,15 +30,15 @@ torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
 torch.backends.cudnn.allow_tf32 = False
 torch.backends.cudnn.deterministic = False
-#torch.backends.cudnn.benchmark = False
-torch.backends.cuda.preferred_blas_library="cublas"
+torch.backends.cudnn.benchmark = False
+#torch.backends.cuda.preferred_blas_library="cublas"
 #torch.backends.cuda.preferred_linalg_library="cusolver"
 
 hftoken = os.getenv("HF_AUTH_TOKEN") 
 
-image_encoder_path = "google/siglip-so400m-patch14-384"
-ipadapter_path = hf_hub_download(repo_id="InstantX/SD3.5-Large-IP-Adapter", filename="ip-adapter.bin")
-model_path = 'ford442/stable-diffusion-3.5-medium-bf16'
+#image_encoder_path = "google/siglip-so400m-patch14-384"
+#ipadapter_path = hf_hub_download(repo_id="InstantX/SD3.5-Large-IP-Adapter", filename="ip-adapter.bin")
+#model_path = 'ford442/stable-diffusion-3.5-medium-bf16'
 
 def upload_to_ftp(filename):
     try:
@@ -62,9 +62,9 @@ checkpoint = "microsoft/Phi-3.5-mini-instruct"
 #vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
 vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #, device_map='cpu') #.to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
 
-pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(device=device, dtype=torch.bfloat16)
+#pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(device=device, dtype=torch.bfloat16)
 #pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(torch.device("cuda:0"))
-#pipe = StableDiffusion3Pipeline.from_pretrained("ford442/RealVis_Medium_1.0b_bf16", torch_dtype=torch.bfloat16)
+pipe = StableDiffusion3Pipeline.from_pretrained("ford442/RealVis_Medium_1.0b_bf16", torch_dtype=torch.bfloat16)
 #pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')
 
 # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
@@ -76,9 +76,9 @@ pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-me
 #pipe = torch.compile(pipe)
 # pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear")
 
-refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16",vae = vaeXL, requires_aesthetics_score=True) #.to(torch.bfloat16)
+#refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16",vae = vaeXL, requires_aesthetics_score=True) #.to(torch.bfloat16)
 #refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float32, requires_aesthetics_score=True, device_map='balanced')
-refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
+#refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
 #refiner.enable_model_cpu_offload()
 
 #pipe.to(device=device, dtype=torch.bfloat16)
@@ -94,7 +94,7 @@ tokenizer.tokenizer_legacy=False
 model = AutoModelForCausalLM.from_pretrained(checkpoint).to('cuda')
 #model = torch.compile(model)
 
-upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
+upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1") #.to(torch.device("cuda:0"))
 
 def filter_text(text,phraseC):
   """Filters out the text up to and including 'Rewritten Prompt:'."""
@@ -136,7 +136,7 @@ def infer(
     latent_file,  # Add latents file input
     progress=gr.Progress(track_tqdm=True),
 ):
-    upscaler_2.to(torch.device('cpu'))
+    #upscaler_2.to(torch.device('cpu'))
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -193,7 +193,7 @@ def infer(
     else:
         enhanced_prompt = prompt
         enhanced_prompt_2 = prompt
-    model.to('cpu')
+    #model.to('cpu')
     if latent_file:  # Check if a latent file is provided
       #  initial_latents = pipe.prepare_latents(
       #      batch_size=1,
@@ -241,6 +241,7 @@ def infer(
             max_sequence_length=512
         ).images[0]
         print('-- got image --')
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         #sd35_image = pipe.vae.decode(sd_image / 0.18215).sample
        # sd35_image = sd35_image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
        # sd35_image = (sd35_image * 255).round().astype("uint8")
@@ -248,7 +249,7 @@ def infer(
       #  sd35_path = f"sd35_{seed}.png"
        # image_pil.save(sd35_path,optimize=False,compress_level=0)
        # upload_to_ftp(sd35_path)
-        sd35_path = f"sd35_{seed}.png"
+        sd35_path = f"sd35_{timestamp}.png"
         sd_image.save(sd35_path,optimize=False,compress_level=0)
         upload_to_ftp(sd35_path)
         # Convert the generated image to a tensor
@@ -261,6 +262,7 @@ def infer(
     #torch.save(generated_latents, latent_path)
     #upload_to_ftp(latent_path)
     #refiner.scheduler.set_timesteps(num_inference_steps,device)
+    '''
     pipe.to(torch.device('cpu'))
     refiner.to(device=device, dtype=torch.bfloat16)
     refine = refiner(
@@ -276,13 +278,14 @@ def infer(
     refine.save(refine_path,optimize=False,compress_level=0)
     upload_to_ftp(refine_path)
     refiner.to(torch.device('cpu'))
+    '''
     upscaler_2.to(torch.device('cuda'))
     with torch.no_grad():
-        upscale2 = upscaler_2(refine, tiling=True, tile_width=256, tile_height=256)
+        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
-    upscaler_2.to(torch.device('cpu'))
+    #upscaler_2.to(torch.device('cpu'))
     downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
-    upscale_path = f"sd35_upscale_{seed}.png"
+    upscale_path = f"sd35_upscale_{timestamp}.png"
     downscale2.save(upscale_path,optimize=False,compress_level=0)
     upload_to_ftp(upscale_path)
     return refine, seed, enhanced_prompt