Update utils.py

utils.py

@@ -20,7 +20,6 @@ from diffusers import StableDiffusionPipeline, DiffusionPipeline
 height, width       = 512, 512
 guidance_scale      = 8
 custom_loss_scale   = 200
-num_inference_steps = 50
 batch_size          = 1
 torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
@@ -68,7 +67,7 @@ def latents_to_pil(latents):
     pil_images = [Image.fromarray(image) for image in images]
     return pil_images
     
-def generate_latents(prompts, seed_nums, loss_apply=False):
+def generate_latents(prompts, num_inference_steps, seed_nums, loss_apply=False):
     
     generator = torch.manual_seed(seed_nums)
     
@@ -141,16 +140,16 @@ def generate_latents(prompts, seed_nums, loss_apply=False):
 def pil_to_np(image):
     return np.array(image)
     
-def generate_gradio_images(prompt):
+def generate_gradio_images(prompt, num_inference_steps):
     # after loss is applied
     latents_list = [] 
     loss_flag = False
     for seed_no, sd in zip(seed_list, sdconcepts):
         prompts = [f'{prompt} {sd}']
-        latents = generate_latents(prompts, seed_no, loss_apply=loss_flag)
+        latents = generate_latents(prompts,num_inference_steps, seed_no, loss_apply=loss_flag)
         latents_list.append(latents)
     loss_flag = True
     for seed_no, sd in zip(seed_list, sdconcepts):
         prompts = [f'{prompt} {sd}']
-        latents = generate_latents(prompts, seed_no, loss_apply=loss_flag)
+        latents = generate_latents(prompts,num_inference_steps, seed_no, loss_apply=loss_flag)
         latents_list.append(latents)