Files changed (1) hide show
  1. app.py +10 -10
app.py CHANGED
@@ -52,17 +52,17 @@ current_model = models[1] if is_colab else models[0]
52
  current_model_path = current_model.path
53
 
54
  if is_colab:
55
- pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
56
 
57
- else: # download all models
58
  print(f"{datetime.datetime.now()} Downloading vae...")
59
- vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16)
60
  for model in models:
61
  try:
62
  print(f"{datetime.datetime.now()} Downloading {model.name} model...")
63
- unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
64
- model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
65
- model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
66
  except Exception as e:
67
  print(f"{datetime.datetime.now()} Failed to load model " + model.name + ": " + str(e))
68
  models.remove(model)
@@ -98,7 +98,7 @@ def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0
98
  current_model = model
99
  model_path = current_model.path
100
 
101
- generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
102
 
103
  try:
104
  if img is not None:
@@ -119,7 +119,7 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
119
  current_model_path = model_path
120
 
121
  if is_colab or current_model == custom_model:
122
- pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
123
  else:
124
  pipe = pipe.to("cpu")
125
  pipe = current_model.pipe_t2i
@@ -152,7 +152,7 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
152
  current_model_path = model_path
153
 
154
  if is_colab or current_model == custom_model:
155
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
156
  else:
157
  pipe = pipe.to("cpu")
158
  pipe = current_model.pipe_i2i
@@ -163,7 +163,7 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
163
 
164
  prompt = current_model.prefix + prompt
165
  ratio = min(height / img.height, width / img.width)
166
- img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
167
  result = pipe(
168
  prompt,
169
  negative_prompt = neg_prompt,
 
52
  current_model_path = current_model.path
53
 
54
  if is_colab:
55
+ pipe = StableDiffusionPipeline.from_pretrained(current_model.path, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
56
 
57
+ else: # download all models All models needed torch_dtype=torch.float16 removed for CPU. Add a If/Else so you can use both
58
  print(f"{datetime.datetime.now()} Downloading vae...")
59
+ vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae")
60
  for model in models:
61
  try:
62
  print(f"{datetime.datetime.now()} Downloading {model.name} model...")
63
+ unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet")
64
+ model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler)
65
+ model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler)
66
  except Exception as e:
67
  print(f"{datetime.datetime.now()} Failed to load model " + model.name + ": " + str(e))
68
  models.remove(model)
 
98
  current_model = model
99
  model_path = current_model.path
100
 
101
+ generator = torch.Generator('cpu').manual_seed(seed) if seed != 0 else None #not using cuda, I can add another if/else if you like
102
 
103
  try:
104
  if img is not None:
 
119
  current_model_path = model_path
120
 
121
  if is_colab or current_model == custom_model:
122
+ pipe = StableDiffusionPipeline.from_pretrained(current_model_path, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
123
  else:
124
  pipe = pipe.to("cpu")
125
  pipe = current_model.pipe_t2i
 
152
  current_model_path = model_path
153
 
154
  if is_colab or current_model == custom_model:
155
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
156
  else:
157
  pipe = pipe.to("cpu")
158
  pipe = current_model.pipe_i2i
 
163
 
164
  prompt = current_model.prefix + prompt
165
  ratio = min(height / img.height, width / img.width)
166
+ img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.Resampling.LANCZOS) #added Resampling to avoid depreciation error
167
  result = pipe(
168
  prompt,
169
  negative_prompt = neg_prompt,