Update sdfile.py
Browse files
sdfile.py
CHANGED
@@ -21,12 +21,12 @@ def get_pipelines( name:PIPELINES, enable_cpu_offload = False, ) -> StableDiffus
|
|
21 |
pipe = None
|
22 |
|
23 |
if name == "txt2img":
|
24 |
-
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
|
25 |
pipe.unet.load_attn_procs("./")
|
26 |
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
|
27 |
elif name == "sketch2img":
|
28 |
-
controlnet = ControlNetModel.from_pretrained("Abhi5ingh/model_dresscode", torch_dtype=torch.float16)
|
29 |
-
pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", controlnet = controlnet, torch_dtype = torch.float16)
|
30 |
pipe.unet.load_attn_procs("./")
|
31 |
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
|
32 |
|
@@ -73,6 +73,7 @@ def generate(
|
|
73 |
else:
|
74 |
raise Exception(
|
75 |
f"Cannot generate image for pipeline {pipeline_name} and {prompt}")
|
|
|
76 |
image = images[0]
|
77 |
|
78 |
os.makedirs("outputs", exist_ok=True)
|
@@ -84,7 +85,5 @@ def generate(
|
|
84 |
)
|
85 |
image.save(f"{filename}.png")
|
86 |
with open(f"{filename}.txt", "w") as f:
|
87 |
-
f.write(f"Prompt: {prompt}\n\nNegative Prompt:{negative_prompt}"
|
88 |
return image
|
89 |
-
|
90 |
-
)
|
|
|
21 |
pipe = None
|
22 |
|
23 |
if name == "txt2img":
|
24 |
+
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16,cache_dir="D:/huggingface/CACHE/")
|
25 |
pipe.unet.load_attn_procs("./")
|
26 |
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
|
27 |
elif name == "sketch2img":
|
28 |
+
controlnet = ControlNetModel.from_pretrained("Abhi5ingh/model_dresscode", torch_dtype=torch.float16,cache_dir="D:/huggingface/CACHE/")
|
29 |
+
pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", controlnet = controlnet, torch_dtype = torch.float16,cache_dir="D:/huggingface/CACHE/")
|
30 |
pipe.unet.load_attn_procs("./")
|
31 |
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
|
32 |
|
|
|
73 |
else:
|
74 |
raise Exception(
|
75 |
f"Cannot generate image for pipeline {pipeline_name} and {prompt}")
|
76 |
+
images = pipe(**kwargs).images
|
77 |
image = images[0]
|
78 |
|
79 |
os.makedirs("outputs", exist_ok=True)
|
|
|
85 |
)
|
86 |
image.save(f"{filename}.png")
|
87 |
with open(f"{filename}.txt", "w") as f:
|
88 |
+
f.write(f"Prompt: {prompt}\n\nNegative Prompt:{negative_prompt}")
|
89 |
return image
|
|
|
|