Update app.py

app.py

@@ -75,9 +75,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=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16"), use_safetensors=True, requires_aesthetics_score=True).to(device).to(torch.bfloat16)
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16", vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16"), requires_aesthetics_score=True).to(device).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, beta_schedule="scaled_linear")
+refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
 #refiner.enable_model_cpu_offload()
 
 #refiner.scheduler.config.requires_aesthetics_score=False
@@ -86,9 +86,9 @@ refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.
 #refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config, beta_schedule="scaled_linear")
 #refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
 
-tokenizer = AutoTokenizer.from_pretrained(checkpoint, add_prefix_space=False).to(torch.device("cuda:0"))
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, add_prefix_space=False, device_map='balanced')
 tokenizer.tokenizer_legacy=False
-model = AutoModelForCausalLM.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='balanced')
 #model = torch.compile(model)
 
 upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))