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
|
|
|
|
|
|