Update app.py
Browse files
app.py
CHANGED
@@ -16,20 +16,21 @@ pipe = DiffusionPipeline.from_pretrained(
|
|
16 |
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
|
17 |
).to(device)
|
18 |
|
19 |
-
generator =
|
20 |
|
21 |
def infer(prompt, init_image):
|
22 |
init_image = Image.open(init_image).convert("RGB")
|
23 |
init_image = init_image.resize((128, 128))
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
|
|
33 |
|
34 |
#with torch.no_grad():
|
35 |
# torch.cuda.empty_cache()
|
|
|
16 |
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
|
17 |
).to(device)
|
18 |
|
19 |
+
generator = torch.Generator("cuda").manual_seed(0)
|
20 |
|
21 |
def infer(prompt, init_image):
|
22 |
init_image = Image.open(init_image).convert("RGB")
|
23 |
init_image = init_image.resize((128, 128))
|
24 |
+
|
25 |
+
with torch.no_grad():
|
26 |
+
res = pipe.train(
|
27 |
+
prompt,
|
28 |
+
init_image,
|
29 |
+
guidance_scale=7.5,
|
30 |
+
num_inference_steps=50,
|
31 |
+
generator=generator,
|
32 |
+
text_embedding_optimization_steps=100,
|
33 |
+
model_fine_tuning_optimization_steps=500)
|
34 |
|
35 |
#with torch.no_grad():
|
36 |
# torch.cuda.empty_cache()
|