Spaces:
Paused
Paused
remove args
Browse files
app.py
CHANGED
@@ -20,12 +20,15 @@ import argparse
|
|
20 |
|
21 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
22 |
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
pretrained_model =
|
|
|
|
|
|
|
29 |
|
30 |
# Check for different hardware architectures
|
31 |
if torch.cuda.is_available():
|
@@ -45,18 +48,18 @@ else:
|
|
45 |
print(f"Using device: {device}")
|
46 |
|
47 |
# Load models
|
48 |
-
if
|
49 |
torch_dtype = torch.float32
|
50 |
-
elif
|
51 |
torch_dtype = torch.float16
|
52 |
-
elif
|
53 |
torch_dtype = torch.bfloat16
|
54 |
else:
|
55 |
-
raise ValueError(f"Invalid precision: {
|
56 |
|
57 |
-
controlnet = ControlNetModel.from_pretrained(
|
58 |
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
59 |
-
|
60 |
)
|
61 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
62 |
pipe = pipe.to(device)
|
|
|
20 |
|
21 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
22 |
|
23 |
+
# parse= argparse.ArgumentParser()
|
24 |
+
# parseadd_argument('--pretrained_model', type=str, default='runwayml/stable-diffusion-v1-5')
|
25 |
+
# parseadd_argument('--controlnet', type=str, default='controlnet')
|
26 |
+
# parseadd_argument('--precision', type=str, default='fp32')
|
27 |
+
# = parseparse_)
|
28 |
+
# pretrained_model = pretrained_model
|
29 |
+
pretrained_model = 'runwayml/stable-diffusion-v1-5'
|
30 |
+
controlnet = 'checkpoint-36000/controlnet'
|
31 |
+
precision = 'bf16'
|
32 |
|
33 |
# Check for different hardware architectures
|
34 |
if torch.cuda.is_available():
|
|
|
48 |
print(f"Using device: {device}")
|
49 |
|
50 |
# Load models
|
51 |
+
if precision == 'fp32':
|
52 |
torch_dtype = torch.float32
|
53 |
+
elif precision == 'fp16':
|
54 |
torch_dtype = torch.float16
|
55 |
+
elif precision == 'bf16':
|
56 |
torch_dtype = torch.bfloat16
|
57 |
else:
|
58 |
+
raise ValueError(f"Invalid precision: {precision}")
|
59 |
|
60 |
+
controlnet = ControlNetModel.from_pretrained(controlnet, torch_dtype=torch_dtype)
|
61 |
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
62 |
+
pretrained_model, controlnet=controlnet, torch_dtype=torch_dtype
|
63 |
)
|
64 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
65 |
pipe = pipe.to(device)
|