Abhi5ingh commited on
Commit
eb7a796
1 Parent(s): 9d157fe

Update sdfile.py

Browse files
Files changed (1) hide show
  1. sdfile.py +5 -6
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