changes

app.py

@@ -1,33 +1,23 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import UNet2DConditionModel, StableDiffusionXLPipeline, EulerDiscreteScheduler
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
+from diffusers import DiffusionPipeline
 import rembg
-from io import BytesIO
-import PIL.Image as Image
-import cv2
-import numpy
 
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-ckpt = "sdxl_lightning_4step_unet.safetensors"
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+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")
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
 
 # Function to generate an image from text using diffusion
 @spaces.GPU
 def generate_image(prompt):
     prompt += "no background, side view, minimalist shot"
     
-    image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
+    image = pipe(prompt).images[0]
+    image2 = rembg.remove(image)
 
-    return image
+    return image, image2
 
 _TITLE = "Shoe Generator"
 with gr.Blocks(_TITLE) as ShoeGen:
@@ -37,8 +27,9 @@ with gr.Blocks(_TITLE) as ShoeGen:
             # neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
             button_gen = gr.Button("Generate Image")
         with gr.Column():
-            image = gr.Image(label="Generated Image", show_download_button=True)   
+            image = gr.Image(label="Generated Image", show_download_button=True) 
+            image_nobg = gr.Image(label="Generated Image (No Background)", show_download_button=True) 
     
-    button_gen.click(generate_image, inputs=[prompt], outputs=[image])
+    button_gen.click(generate_image, inputs=[prompt], outputs=[image, image_nobg])
 
 ShoeGen.launch()