Update app.py

app.py

@@ -6,11 +6,14 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
+from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, DPMSolverMultistepScheduler
 from huggingface_hub import hf_hub_download, InferenceClient
 
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+refiner.to("cuda")
+
 pipe_fast = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V4.0_Lightning", torch_dtype=torch.float16, vae=vae, use_safetensors=True)
 pipe_fast.to("cuda")
 
@@ -53,7 +56,6 @@ def promptifier(prompt):
     return "".join([response.token.text for response in stream if response.token.text != "</s>"])
 
 client_image = InferenceClient("stabilityai/stable-diffusion-3-medium-diffusers")
-refiner = InferenceClient("stabilityai/stable-diffusion-xl-refiner-1.0")
 # Generator
 @spaces.GPU(duration=60, queue=False)
 def king(type ,
@@ -80,13 +82,14 @@ def king(type ,
             width = input_image.width, height = input_image.height,
             num_inference_steps=steps, generator=generator, output_type="latent",
         ).images
-        refine = refiner.image_to_image(
+        refine = refiner(
             prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
             negative_prompt = negative_prompt,
             guidance_scale=7.5,
-            num_inference_steps=int(steps/3),
-            image=output_image
-        )
+            num_inference_steps=steps,
+            image=output_image,
+            generator=generator,
+        ).images[0]  
         return seed, refine
     else :
         if randomize_seed:
@@ -110,10 +113,12 @@ def king(type ,
             guidance_scale = guidance_scale, 
             num_inference_steps = steps, 
             width = width, height = height )
-            refine = refiner.image_to_image( 
-                prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
-                negative_prompt = negative_prompt, guidance_scale = 7.5,
-                num_inference_steps= int(steps/3), image=image )     
+            refine = refiner( prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
+                    negative_prompt = negative_prompt,
+                    guidance_scale = 7.5,
+                    num_inference_steps=  steps,
+                    image=image, generator=generator,
+            ).images[0]        
         return seed, refine
 
 client = InferenceClient()