Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
@@ -42,19 +42,18 @@ last_mode = "txt2img"
|
|
42 |
current_model = models[0]
|
43 |
current_model_path = current_model.path
|
44 |
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
pipe = models[0].pipe_t2i
|
58 |
|
59 |
if torch.cuda.is_available():
|
60 |
pipe = pipe.to("cuda")
|
|
|
42 |
current_model = models[0]
|
43 |
current_model_path = current_model.path
|
44 |
|
45 |
+
print(f"{datetime.datetime.now()} Downloading vae...")
|
46 |
+
vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae")
|
47 |
+
for model in models:
|
48 |
+
try:
|
49 |
+
print(f"{datetime.datetime.now()} Downloading {model.name} model...")
|
50 |
+
unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet")
|
51 |
+
model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler)
|
52 |
+
model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler)
|
53 |
+
except Exception as e:
|
54 |
+
print(f"{datetime.datetime.now()} Failed to load model " + model.name + ": " + str(e))
|
55 |
+
models.remove(model)
|
56 |
+
pipe = models[0].pipe_t2i
|
|
|
57 |
|
58 |
if torch.cuda.is_available():
|
59 |
pipe = pipe.to("cuda")
|