refactor: Update image generation pipeline to use playground v2.5

launch/image_generation.py

@@ -4,19 +4,16 @@ import gradio as gr
 import rembg
 import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
+from diffusers import DiffusionPipeline
 
 from src.utils.infer_util import (remove_background, resize_foreground)
 
 
-# Load StableDiffusionXL model
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+pipe = DiffusionPipeline.from_pretrained(
+    "playgroundai/playground-v2.5-1024px-aesthetic",
+    torch_dtype=torch.float16,
+    variant="fp16"
+).to("cuda")
 
 
 def generate_prompt(subject, style, color_scheme, angle, lighting_type, additional_details):
@@ -25,18 +22,9 @@ def generate_prompt(subject, style, color_scheme, angle, lighting_type, addition
 
 @spaces.GPU
 def generate_image(subject, style, color_scheme, angle, lighting_type, additional_details):
-    checkpoint = "sdxl_lightning_8step_unet.safetensors"
-    num_inference_steps = 8
-
-    pipe.scheduler = EulerDiscreteScheduler.from_config(
-        pipe.scheduler.config, timestep_spacing="trailing")
-    pipe.unet.load_state_dict(
-        load_file(hf_hub_download(repo, checkpoint), device="cuda"))
-
     prompt = generate_prompt(subject, style, color_scheme,
                              angle, lighting_type, additional_details)
-    results = pipe(
-        prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
+    results = pipe(prompt, num_inference_steps=25, guidance_scale=7.5)
     return results.images[0]
 
 

launch/utils.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 
+
 def find_cuda():
     cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
     if cuda_home and os.path.exists(cuda_home):
@@ -11,4 +12,4 @@ def find_cuda():
         cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
         return cuda_path
 
-    return None
+    return None