Yardenfren commited on
Commit
2ad93a6
1 Parent(s): fa662f9

Update inf.py

Browse files
Files changed (1) hide show
  1. inf.py +11 -7
inf.py CHANGED
@@ -18,15 +18,19 @@ class InferencePipeline:
18
  self.hf_token = hf_token
19
  self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
20
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
21
- if self.device.type == 'cpu':
22
- self.pipe = StableDiffusionXLPipeline.from_pretrained(
23
- self.base_model_id, use_auth_token=self.hf_token)
24
- else:
25
- self.pipe = StableDiffusionXLPipeline.from_pretrained(
 
 
 
 
 
26
  self.base_model_id,
27
- # torch_dtype=torch.float16,
28
  use_auth_token=self.hf_token)
29
- self.pipe = self.pipe.to(self.device)
30
  self.content_lora_model_id = None
31
  self.style_lora_model_id = None
32
 
 
18
  self.hf_token = hf_token
19
  self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
20
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
21
+ # if self.device.type == 'cpu':
22
+ # self.pipe = StableDiffusionXLPipeline.from_pretrained(
23
+ # self.base_model_id, use_auth_token=self.hf_token)
24
+ # else:
25
+ # self.pipe = StableDiffusionXLPipeline.from_pretrained(
26
+ # self.base_model_id,
27
+ # torch_dtype=torch.float16,
28
+ # use_auth_token=self.hf_token)
29
+ # self.pipe = self.pipe.to(self.device)
30
+ self.pipe = StableDiffusionXLPipeline.from_pretrained(
31
  self.base_model_id,
32
+ torch_dtype=torch.float16,
33
  use_auth_token=self.hf_token)
 
34
  self.content_lora_model_id = None
35
  self.style_lora_model_id = None
36