~~

app.py

@@ -57,6 +57,7 @@ def load_pipeline(model_name):
         model_name,
         torch_dtype=torch.float16,
         custom_pipeline="lpw_stable_diffusion_xl",
+        safety_checker = None,
         use_safetensors=True,
         add_watermarker=False,
         use_auth_token=HF_TOKEN,
@@ -64,3 +65,50 @@ def load_pipeline(model_name):
     pipe.to(device)
     return pipe
 
+@spaces.GPU
+def generate(
+    prompt: str,
+    negative_prompt: str = None,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 5.0,
+    num_inference_steps: int = 24,
+    sampler: str = "Euler a",
+    clip_skip: int = 1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    
+    generator = seed_everything(seed)
+    pipe.scheduler = get_scheduler(pipe.scheduler.config, sampler)
+    pipe.text_encoder = CLIPTextModel.from_pretrained(
+        MODEL,
+        subfolder = "text_encoder",
+        num_hidden_layers = 12 - (clip_skip - 1),
+        torch_dtype = torch.float16
+    )
+
+    try:
+
+        img = pipe(
+            prompt = prompt,
+            negative_prompt = negative_prompt,
+            width = width,
+            height = height,
+            guidance_scale = guidance_scale,
+            num_inference_steps = num_inference_steps,
+            generator = generator,
+            output_type="pil",
+        ).images[0]
+
+        return img
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+if torch.cuda.is_available():
+    pipe = load_pipeline(MODEL)
+    print("Loaded on Device!")
+else:
+    pipe = None
+