Spaces:
Runtime error
Runtime error
Jv
commited on
Commit
•
66dc286
1
Parent(s):
20381f1
Update app.py
Browse files
app.py
CHANGED
@@ -18,7 +18,7 @@ checkpoints = {
|
|
18 |
"4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
|
19 |
"8-Step" : ["sdxl_lightning_8step_unet.safetensors", 8],
|
20 |
}
|
21 |
-
|
22 |
|
23 |
# Ensure model and scheduler are initialized in GPU-enabled function
|
24 |
if torch.cuda.is_available():
|
@@ -49,17 +49,18 @@ if SAFETY_CHECKER:
|
|
49 |
# Function
|
50 |
@spaces.GPU(enable_queue=True)
|
51 |
def generate_image(prompt, ckpt):
|
52 |
-
|
53 |
-
print(prompt, ckpt)
|
54 |
-
|
55 |
checkpoint = checkpoints[ckpt][0]
|
56 |
num_inference_steps = checkpoints[ckpt][1]
|
57 |
|
58 |
-
if
|
59 |
-
|
60 |
-
pipe.
|
61 |
-
|
|
|
|
|
62 |
|
|
|
63 |
results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
|
64 |
|
65 |
if SAFETY_CHECKER:
|
|
|
18 |
"4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
|
19 |
"8-Step" : ["sdxl_lightning_8step_unet.safetensors", 8],
|
20 |
}
|
21 |
+
|
22 |
|
23 |
# Ensure model and scheduler are initialized in GPU-enabled function
|
24 |
if torch.cuda.is_available():
|
|
|
49 |
# Function
|
50 |
@spaces.GPU(enable_queue=True)
|
51 |
def generate_image(prompt, ckpt):
|
52 |
+
|
|
|
|
|
53 |
checkpoint = checkpoints[ckpt][0]
|
54 |
num_inference_steps = checkpoints[ckpt][1]
|
55 |
|
56 |
+
if num_inference_steps==1:
|
57 |
+
# Ensure sampler uses "trailing" timesteps and "sample" prediction type for 1-step inference.
|
58 |
+
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
|
59 |
+
else:
|
60 |
+
# Ensure sampler uses "trailing" timesteps.
|
61 |
+
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
62 |
|
63 |
+
pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
|
64 |
results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
|
65 |
|
66 |
if SAFETY_CHECKER:
|