CHANGES

app.py

@@ -1,34 +1,34 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, AutoencoderKL
+from diffusers import UNet2DConditionModel, StableDiffusionXLPipeline, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 import rembg
 from io import BytesIO
 import PIL.Image as Image
 import cv2
 import numpy
 
-model_id = "dataautogpt3/OpenDalleV1.1"
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-pipe = DiffusionPipeline.from_pretrained(model_id,
-                                         vae=vae,
-                                         torch_dtype=torch.float16,
-                                        use_safetensors=True,
-                                        variant="fp16")
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_4step_unet.safetensors"
 
-pipe.to("cuda")
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
 # Function to generate an image from text using diffusion
 @spaces.GPU
 def generate_image(prompt, neg_prompt):
     prompt += "no background, side view, minimalist shot"
-    image_bytes = pipe(prompt, negative_prompt=neg_prompt)
-    print(image_bytes)
-    pil_image = Image.open(BytesIO(image_bytes))
-    opencvImage = cv2.cvtColor(numpy.array(pil_image), cv2.COLOR_RGB2BGR)
+    
+    image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
+    image2  = rembg.remove(image)
 
-    cv2_image = rembg.remove(opencvImage)
-    return pil_image, cv2_image
+    return image, image2
 
 _TITLE = "Shoe Generator"
 with gr.Blocks(_TITLE) as ShoeGen:

requirements.txt

@@ -8,4 +8,5 @@ rembg
 Pillow
 Python-IO
 numpy
-opencv-python
+opencv-python
+huggingface-hub