Update app.py

app.py

@@ -6,14 +6,11 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
 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")
 
@@ -56,6 +53,7 @@ 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 ,
@@ -82,14 +80,13 @@ def king(type ,
             width = input_image.width, height = input_image.height,
             num_inference_steps=steps, generator=generator, output_type="latent",
         ).images
-        refine = refiner(
+        refine = refiner.image_to_image(
             prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
             negative_prompt = negative_prompt,
             guidance_scale=7.5,
-            num_inference_steps=steps,
-            image=output_image,
-            generator=generator,
-        ).images[0]  
+            num_inference_steps=int(steps/3),
+            image=output_image
+        )
         return seed, refine
     else :
         if randomize_seed:
@@ -113,12 +110,10 @@ def king(type ,
             guidance_scale = guidance_scale, 
             num_inference_steps = steps, 
             width = width, height = height )
-            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]        
+            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 )     
         return seed, refine
 
 client = InferenceClient()