Allow large pictures

app.py

@@ -9,7 +9,7 @@ device = "cuda"
 
 base = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
-).to(device)
+)
 
 refiner = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-refiner-1.0",
@@ -18,7 +18,7 @@ refiner = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     torch_dtype=torch.float16,
     use_safetensors=True,
     variant="fp16",
-).to(device)
+)
 
 base.scheduler = DPMSolverMultistepScheduler.from_config(base.scheduler.config)
 refiner.scheduler = DPMSolverMultistepScheduler.from_config(base.scheduler.config)
@@ -46,17 +46,20 @@ def inference(image, map, gs, prompt, negative_prompt):
     validate_inputs(image, map)
     image = preprocess_image(image)
     map = preprocess_map(map)
-    edited_images = base(prompt=prompt, original_image=image, image=image, strength=1, guidance_scale=gs,
+    base_cuda = base.to(device)
+    edited_images = base_cuda(prompt=prompt, original_image=image, image=image, strength=1, guidance_scale=gs,
                          num_images_per_prompt=1,
                          negative_prompt=negative_prompt,
                          map=map,
                          num_inference_steps=NUM_INFERENCE_STEPS, denoising_end=0.8, output_type="latent").images
-
-    edited_images = refiner(prompt=prompt, original_image=image, image=edited_images, strength=1, guidance_scale=7.5,
+    base_cuda=None
+    refiner_cuda = refiner.to(device)
+    edited_images = refiner_cuda(prompt=prompt, original_image=image, image=edited_images, strength=1, guidance_scale=7.5,
                             num_images_per_prompt=1,
                             negative_prompt=negative_prompt,
                             map=map,
                             num_inference_steps=NUM_INFERENCE_STEPS, denoising_start=0.8).images[0]
+    refiner_cuda=None
     return edited_images