Update app.py

app.py

@@ -71,7 +71,6 @@ def gaussian_kernel(kernel_size=3, sigma=1.0, channels=3):
     gaussian_1d = gaussian_1d / gaussian_1d.sum()
     gaussian_2d = gaussian_1d[:, None] * gaussian_1d[None, :]
     kernel = gaussian_2d[None, None, :, :].repeat(channels, 1, 1, 1)
-    
     return kernel
 
 def gaussian_filter(latents, kernel_size=3, sigma=1.0):
@@ -88,9 +87,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     """
     std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
     std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-    # rescale the results from guidance (fixes overexposure)
     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
     return noise_cfg
 
@@ -144,7 +141,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         add_watermarker: Optional[bool] = None,
     ):
         super().__init__()
-
         self.register_modules(
             vae=vae,
             text_encoder=text_encoder,
@@ -158,27 +154,21 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
         self.default_sample_size = self.unet.config.sample_size
-
         add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
-
         if add_watermarker:
             self.watermark = StableDiffusionXLWatermarker()
         else:
             self.watermark = None
 
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
     def enable_vae_slicing(self):
         self.vae.enable_slicing()
 
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
     def disable_vae_slicing(self):
         self.vae.disable_slicing()
 
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling
     def enable_vae_tiling(self):
         self.vae.enable_tiling()
 
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling
     def disable_vae_tiling(self):
         self.vae.disable_tiling()
 
@@ -198,25 +188,20 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         lora_scale: Optional[float] = None,
     ):
         device = device or self._execution_device
-
-        # set lora scale so that monkey patched LoRA function of text encoder can correctly access it
         if lora_scale is not None and isinstance(self, LoraLoaderMixin):
             self._lora_scale = lora_scale
             adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
             adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
-
         if prompt is not None and isinstance(prompt, str):
             batch_size = 1
         elif prompt is not None and isinstance(prompt, list):
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
-
         tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
         text_encoders = (
             [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
         )
-
         if prompt_embeds is None:
             prompt_2 = prompt_2 or prompt
             prompt_embeds_list = []
@@ -224,7 +209,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
                 if isinstance(self, TextualInversionLoaderMixin):
                     prompt = self.maybe_convert_prompt(prompt, tokenizer)
-
                 text_inputs = tokenizer(
                     prompt,
                     padding="max_length",
@@ -232,10 +216,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                     truncation=True,
                     return_tensors="pt",
                 )
-
                 text_input_ids = text_inputs.input_ids
                 untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
                 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
                     text_input_ids, untruncated_ids
                 ):
@@ -244,19 +226,14 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                         "The following part of your input was truncated because CLIP can only handle sequences up to"
                         f" {tokenizer.model_max_length} tokens: {removed_text}"
                     )
-
                 prompt_embeds = text_encoder(
                     text_input_ids.to(device),
                     output_hidden_states=True,
                 )
-
                 pooled_prompt_embeds = prompt_embeds[0]
                 prompt_embeds = prompt_embeds.hidden_states[-2]
-
                 prompt_embeds_list.append(prompt_embeds)
-
             prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
         zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
         if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
             negative_prompt_embeds = torch.zeros_like(prompt_embeds)
@@ -264,7 +241,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         elif do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
             negative_prompt_2 = negative_prompt_2 or negative_prompt
-
             if prompt is not None and type(prompt) is not type(negative_prompt):
                 raise TypeError(
                     f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
@@ -280,12 +256,10 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 )
             else:
                 uncond_tokens = [negative_prompt, negative_prompt_2]
-
             negative_prompt_embeds_list = []
             for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
                 if isinstance(self, TextualInversionLoaderMixin):
                     negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
-
                 max_length = prompt_embeds.shape[1]
                 uncond_input = tokenizer(
                     negative_prompt,
@@ -294,29 +268,23 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                     truncation=True,
                     return_tensors="pt",
                 )
-
                 negative_prompt_embeds = text_encoder(
                     uncond_input.input_ids.to(device),
                     output_hidden_states=True,
                 )
                 negative_pooled_prompt_embeds = negative_prompt_embeds[0]
                 negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
-
                 negative_prompt_embeds_list.append(negative_prompt_embeds)
-
             negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
-
         prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
         bs_embed, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
         prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
         if do_classifier_free_guidance:
             seq_len = negative_prompt_embeds.shape[1]
             negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
             negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
             negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
         pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
             bs_embed * num_images_per_prompt, -1
         )
@@ -324,7 +292,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
                 bs_embed * num_images_per_prompt, -1
             )
-
         return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
 
     def prepare_extra_step_kwargs(self, generator, eta):
@@ -332,7 +299,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         extra_step_kwargs = {}
         if accepts_eta:
             extra_step_kwargs["eta"] = eta
-
         accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
         if accepts_generator:
             extra_step_kwargs["generator"] = generator
@@ -355,14 +321,12 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
     ):
         if height % 8 != 0 or width % 8 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
         if (callback_steps is None) or (
             callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
         ):
             raise ValueError(
                 f"`callback_steps` has to be a positive integer but is {callback_steps} of type {type(callback_steps)}."
             )
-
         if prompt is not None and prompt_embeds is not None:
             raise ValueError(
                 f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to only forward one of the two."
@@ -379,7 +343,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
         elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
             raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-
         if negative_prompt is not None and negative_prompt_embeds is not None:
             raise ValueError(
                 f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to only forward one of the two."
@@ -388,26 +351,21 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             raise ValueError(
                 f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to only forward one of the two."
             )
-
         if prompt_embeds is not None and negative_prompt_embeds is not None:
             if prompt_embeds.shape != negative_prompt_embeds.shape:
                 raise ValueError(
                     "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds` {negative_prompt_embeds.shape}."
                 )
-
         if prompt_embeds is not None and pooled_prompt_embeds is None:
             raise ValueError(
                 "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
             )
-
         if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
             raise ValueError(
                 "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
             )
-
         if max(height, width) % 1024 != 0:
             raise ValueError(f"the larger one of `height` and `width` has to be divisible by 1024 but are {height} and {width}.")
-
         if num_images_per_prompt != 1:
             warnings.warn("num_images_per_prompt != 1 is not supported by AccDiffusion and will be ignored.")
             num_images_per_prompt = 1
@@ -418,29 +376,24 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             raise ValueError(
                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch size of {batch_size}. Make sure the batch size matches the length of the generators."
             )
-
         if latents is None:
             latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         else:
             latents = latents.to(device)
-
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
     def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
         add_time_ids = list(original_size + crops_coords_top_left + target_size)
-
         passed_add_embed_dim = (
             self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim
         )
         expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
-
         if expected_add_embed_dim != passed_add_embed_dim:
             raise ValueError(
                 f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. \
                 The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
             )
-
         add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
         return add_time_ids
 
@@ -456,7 +409,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             h_end = h_start + window_size
             w_start = int((i % num_blocks_width) * stride)
             w_end = w_start + window_size
-
             if h_end > height:
                 h_start = int(h_start + height - h_end)
                 h_end = int(height)
@@ -469,7 +421,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             if w_start < 0:
                 w_end = int(w_end - w_start)
                 w_start = 0
-
             if random_jitter:
                 jitter_range = (window_size - stride) // 4
                 w_jitter = 0
@@ -480,7 +431,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                     w_jitter = random.randint(-jitter_range, 0)
                 elif (w_start != 0) and (w_end == width):
                     w_jitter = random.randint(0, jitter_range)
-
                 if (h_start != 0) and (h_end != height):
                     h_jitter = random.randint(-jitter_range, jitter_range)
                 elif (h_start == 0) and (h_end != height):
@@ -491,7 +441,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 h_end = h_end + h_jitter + jitter_range
                 w_start = w_start + w_jitter + jitter_range
                 w_end = w_end + w_jitter + jitter_range
-            
             views.append((h_start, h_end, w_start, w_end))
         return views
 
@@ -527,7 +476,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 continue
             cross_att_count += 1
             attn_procs[name] = P2PCrossAttnProcessor(controller=controller, place_in_unet=place_in_unet)
-
         self.unet.set_attn_processor(attn_procs)
         controller.num_att_layers = cross_att_count
         return ori_attn_processors
@@ -540,7 +488,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module
         else:
             raise ImportError("Offloading requires `accelerate v0.17.0` or higher.")
-
         is_model_cpu_offload = False
         is_sequential_cpu_offload = False
         recursive = False
@@ -560,7 +507,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             **kwargs,
         )
         self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet)
-
         text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
         if len(text_encoder_state_dict) > 0:
             self.load_lora_into_text_encoder(
@@ -570,7 +516,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 prefix="text_encoder",
                 lora_scale=self.lora_scale,
             )
-
         text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k}
         if len(text_encoder_2_state_dict) > 0:
             self.load_lora_into_text_encoder(
@@ -580,7 +525,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 prefix="text_encoder_2",
                 lora_scale=self.lora_scale,
             )
-
         if is_model_cpu_offload:
             self.enable_model_cpu_offload()
         elif is_sequential_cpu_offload:
@@ -599,24 +543,19 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
         safe_serialization: bool = True,
     ):
         state_dict = {}
-
         def pack_weights(layers, prefix):
             layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
             layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
             return layers_state_dict
-
         if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
             raise ValueError(
                 "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers` or `text_encoder_2_lora_layers`."
             )
-
         if unet_lora_layers:
             state_dict.update(pack_weights(unet_lora_layers, "unet"))
-
         if text_encoder_lora_layers and text_encoder_2_lora_layers:
             state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder"))
             state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
-
         self.write_lora_layers(
             state_dict=state_dict,
             save_directory=save_directory,
@@ -697,39 +636,29 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
 
         Examples:
         """
-    
         if debug:
             num_inference_steps = 1
-        
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor
-
         x1_size = self.default_sample_size * self.vae_scale_factor
-
         height_scale = height / x1_size
         width_scale = width / x1_size
         scale_num = int(max(height_scale, width_scale))
         aspect_ratio = min(height_scale, width_scale) / max(height_scale, width_scale)
-
         original_size = original_size or (height, width)
         target_size = target_size or (height, width)
-
         if attn_res is None:
-            attn_res = int(np.ceil(self.default_sample_size * self.vae_scale_factor / 32)), int(np.ceil(self.default_sample_size * self.vae_scale_factor / 32))
+            attn_res = (int(np.ceil(self.default_sample_size * self.vae_scale_factor / 32)), int(np.ceil(self.default_sample_size * self.vae_scale_factor / 32)))
         self.attn_res = attn_res
-
         if lowvram:
             attention_map_device = torch.device("cpu")
         else:
             attention_map_device = self.device
-
         self.controller = create_controller(
             prompt, cross_attention_kwargs, num_inference_steps, tokenizer=self.tokenizer, device=attention_map_device, attn_res=self.attn_res
         )
-
         if save_attention_map or use_md_prompt:
             ori_attn_processors = self.register_attention_control(self.controller)
-
         self.check_inputs(
             prompt,
             prompt_2,
@@ -744,14 +673,12 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             negative_pooled_prompt_embeds,
             num_images_per_prompt,
         )
-
         if prompt is not None and isinstance(prompt, str):
             batch_size = 1
         elif prompt is not None and isinstance(prompt, list):
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
-
         device = self._execution_device
         self.lowvram = lowvram
         if self.lowvram:
@@ -759,13 +686,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             self.unet.cpu()
             self.text_encoder.to(device)
             self.text_encoder_2.to(device)
-
         do_classifier_free_guidance = guidance_scale > 1.0
-
-        text_encoder_lora_scale = (
-            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
-        )
-
+        text_encoder_lora_scale = (cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None)
         (
             prompt_embeds,
             negative_prompt_embeds,
@@ -785,10 +707,8 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
             lora_scale=text_encoder_lora_scale,
         )
-
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
-
         num_channels_latents = self.unet.config.in_channels
         latents = self.prepare_latents(
             batch_size * num_images_per_prompt,
@@ -800,15 +720,11 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             generator,
             latents,
         )
-
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
         add_text_embeds = pooled_prompt_embeds
-
         add_time_ids = self._get_add_time_ids(
             original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
         )
-
         if negative_original_size is not None and negative_target_size is not None:
             negative_add_time_ids = self._get_add_time_ids(
                 negative_original_size,
@@ -818,16 +734,12 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             )
         else:
             negative_add_time_ids = add_time_ids
-
         if do_classifier_free_guidance:
             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0).to(device)
             add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0).to(device)
             add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0).to(device).repeat(batch_size * num_images_per_prompt, 1)
-
         del negative_prompt_embeds, negative_pooled_prompt_embeds, negative_add_time_ids
-
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-
         if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1:
             discrete_timestep_cutoff = int(
                 round(
@@ -836,35 +748,26 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             )
             num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
             timesteps = timesteps[:num_inference_steps]
-
         output_images = []
-        
         ###################################################### Phase Initialization ########################################################
-
         if self.lowvram:
             self.text_encoder.cpu()
             self.text_encoder_2.cpu()
-
         if image_lr == None:
             print("### Phase 1 Denoising ###")
             with self.progress_bar(total=num_inference_steps) as progress_bar:
                 for i, t in enumerate(timesteps):
-    
                     if self.lowvram:
                         self.vae.cpu()
                         self.unet.to(device)
-    
                     latents_for_view = latents
-    
                     latent_model_input = (
                         latents.repeat_interleave(2, dim=0)
                         if do_classifier_free_guidance
                         else latents
                     )
                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
                     added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-
                     noise_pred = self.unet(
                         latent_model_input,
                         t,
@@ -872,25 +775,19 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                         added_cond_kwargs=added_cond_kwargs,
                         return_dict=False,
                     )[0]
-    
                     if do_classifier_free_guidance:
                         noise_pred_uncond, noise_pred_text = noise_pred[::2], noise_pred[1::2]
                         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-    
                     if do_classifier_free_guidance and guidance_rescale > 0.0:
                         noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-    
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-    
                     if t == 1 and use_md_prompt:
                         md_prompts, views_attention = get_multidiffusion_prompts(tokenizer=self.tokenizer, prompts=[prompt], threthod=c, attention_store=self.controller, height=height//scale_num, width=width//scale_num, from_where=["up","down"], random_jitter=True, scale_num=scale_num)
-    
                     if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                         progress_bar.update()
                         if callback is not None and i % callback_steps == 0:
                             step_idx = i // getattr(self.scheduler, "order", 1)
                             callback(step_idx, t, latents)
-
             del latents_for_view, latent_model_input, noise_pred, noise_pred_text, noise_pred_uncond
             if use_md_prompt or save_attention_map:
                 self.recover_attention_control(ori_attn_processors=ori_attn_processors)
@@ -900,7 +797,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             print("### Encoding Real Image ###")
             latents = self.vae.encode(image_lr)
             latents = latents.latent_dist.sample() * self.vae.config.scaling_factor
-
         anchor_mean = latents.mean()
         anchor_std = latents.std()
         if self.lowvram:
@@ -908,12 +804,10 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
             torch.cuda.empty_cache()
         if not output_type == "latent":
             needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
-            
             if self.lowvram:
                 needs_upcasting = False
                 self.unet.cpu()
                 self.vae.to(device)
-
             if needs_upcasting:
                 self.upcast_vae()
                 latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
@@ -924,15 +818,12 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
             if needs_upcasting:
                 self.vae.to(dtype=torch.float16)
-    
         image = self.image_processor.postprocess(image, output_type=output_type)
         if not os.path.exists(f'{result_path}'):
             os.makedirs(f'{result_path}')
-
         image_lr_save_path = f'{result_path}/{image[0].size[0]}_{image[0].size[1]}.png'
         image[0].save(image_lr_save_path)
         output_images.append(image[0])
-
         ####################################################### Phase Upscaling #####################################################
         if use_progressive_upscaling:
             if image_lr == None:
@@ -941,26 +832,21 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 starting_scale = 1
         else:
             starting_scale = scale_num
-
         for current_scale_num in range(starting_scale, scale_num + 1):
             if self.lowvram:
                 latents = latents.to(device)
                 self.unet.to(device)
                 torch.cuda.empty_cache()
-            
             current_height = self.unet.config.sample_size * self.vae_scale_factor * current_scale_num
             current_width = self.unet.config.sample_size * self.vae_scale_factor * current_scale_num
-
             if height > width:
                 current_width = int(current_width * aspect_ratio)
             else:
                 current_height = int(current_height * aspect_ratio)
-
             if upscale_mode == "bicubic_latent" or debug:
                 latents = F.interpolate(latents.to(device), size=(int(current_height / self.vae_scale_factor), int(current_width / self.vae_scale_factor)), mode='bicubic')
             else:
                 raise NotImplementedError
-            
             print("### Phase {} Denoising ###".format(current_scale_num))
             noise_latents = []
             noise = torch.randn_like(latents)
@@ -968,7 +854,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                 noise_latent = self.scheduler.add_noise(latents, noise, timestep.unsqueeze(0))
                 noise_latents.append(noise_latent)
             latents = noise_latents[0]
-
             with self.progress_bar(total=num_inference_steps) as progress_bar:
                 for i, t in enumerate(timesteps):
                     count = torch.zeros_like(latents)
@@ -977,7 +862,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                     if use_skip_residual:
                         c1 = cosine_factor ** cosine_scale_1
                         latents = latents * (1 - c1) + noise_latents[i] * c1
-                    
                     if use_multidiffusion:
                         if use_md_prompt:
                             md_prompt_embeds_list = []
@@ -1005,46 +889,33 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                                 md_prompt_embeds_list.append(torch.cat([md_negative_prompt_embeds, md_prompt_embeds], dim=0).to(device))
                                 md_add_text_embeds_list.append(torch.cat([md_negative_pooled_prompt_embeds, md_pooled_prompt_embeds], dim=0).to(device))
                                 del md_negative_prompt_embeds, md_negative_pooled_prompt_embeds
-
                         if use_md_prompt:
                             random_jitter = True
                             views = [(h_start*4, h_end*4, w_start*4, w_end*4) for h_start, h_end, w_start, w_end in views_attention[current_scale_num]]
                         else:
                             random_jitter = True
                             views = self.get_views(current_height, current_width, stride=stride, window_size=self.unet.config.sample_size, random_jitter=random_jitter)
-
                         views_batch = [views[i : i + view_batch_size] for i in range(0, len(views), view_batch_size)]
-
                         if use_md_prompt:
                             views_prompt_embeds_input = [md_prompt_embeds_list[i : i + view_batch_size] for i in range(0, len(views), view_batch_size)]
                             views_add_text_embeds_input = [md_add_text_embeds_list[i : i + view_batch_size] for i in range(0, len(views), view_batch_size)]
-
                         if random_jitter:
                             jitter_range = int((self.unet.config.sample_size - stride) // 4)
                             latents_ = F.pad(latents, (jitter_range, jitter_range, jitter_range, jitter_range), 'constant', 0)
                         else:
                             latents_ = latents
-
                         count_local = torch.zeros_like(latents_)
                         value_local = torch.zeros_like(latents_)
-                        
                         for j, batch_view in enumerate(views_batch):
                             vb_size = len(batch_view)
                             latents_for_view = torch.cat(
-                                [
-                                    latents_[:, :, h_start:h_end, w_start:w_end]
-                                    for h_start, h_end, w_start, w_end in batch_view
-                                ]
+                                [latents_[:, :, h_start:h_end, w_start:w_end] for h_start, h_end, w_start, w_end in batch_view]
                             )
-
                             latent_model_input = latents_for_view
-                            latent_model_input = (
-                                latent_model_input.repeat_interleave(2, dim=0)
-                                if do_classifier_free_guidance
-                                else latent_model_input
-                            )
+                            latent_model_input = (latent_model_input.repeat_interleave(2, dim=0)
+                                                    if do_classifier_free_guidance
+                                                    else latent_model_input)
                             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-                            
                             add_time_ids_input = []
                             for h_start, h_end, w_start, w_end in batch_view:
                                 add_time_ids_ = add_time_ids.clone()
@@ -1052,7 +923,6 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                                 add_time_ids_[:, 3] = w_start * self.vae_scale_factor
                                 add_time_ids_input.append(add_time_ids_)
                             add_time_ids_input = torch.cat(add_time_ids_input)
-
                             if not use_md_prompt:
                                 prompt_embeds_input = torch.cat([prompt_embeds] * vb_size)
                                 add_text_embeds_input = torch.cat([add_text_embeds] * vb_size)
@@ -1075,30 +945,20 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                                     added_cond_kwargs=md_added_cond_kwargs,
                                     return_dict=False,
                                 )[0]
-
                             if do_classifier_free_guidance:
                                 noise_pred_uncond, noise_pred_text = noise_pred[::2], noise_pred[1::2]
                                 noise_pred = noise_pred_uncond + multi_guidance_scale * (noise_pred_text - noise_pred_uncond)
-
                             if do_classifier_free_guidance and guidance_rescale > 0.0:
                                 noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
                             self.scheduler._init_step_index(t)
-                            latents_denoised_batch = self.scheduler.step(
-                                noise_pred, t, latents_for_view, **extra_step_kwargs, return_dict=False)[0]
-                            
-                            for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip(
-                                latents_denoised_batch.chunk(vb_size), batch_view
-                            ):
+                            latents_denoised_batch = self.scheduler.step(noise_pred, t, latents_for_view, **extra_step_kwargs, return_dict=False)[0]
+                            for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip(latents_denoised_batch.chunk(vb_size), batch_view):
                                 value_local[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised
                                 count_local[:, :, h_start:h_end, w_start:w_end] += 1
-
                         if random_jitter:
                             value_local = value_local[:, :, jitter_range:jitter_range + current_height // self.vae_scale_factor, jitter_range:jitter_range + current_width // self.vae_scale_factor]
                             count_local = count_local[:, :, jitter_range:jitter_range + current_height // self.vae_scale_factor, jitter_range:jitter_range + current_width // self.vae_scale_factor]
-
                         noise_index = i + 1 if i != (len(timesteps) - 1) else i
-
                         value_local = torch.where(count_local == 0, noise_latents[noise_index], value_local)
                         count_local = torch.where(count_local == 0, torch.ones_like(count_local), count_local)
                         if use_dilated_sampling:
@@ -1107,19 +967,15 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                             count += torch.ones_like(value_local) * (1 - c2)
                         else:
                             value += value_local / count_local
-                            count += torch.ones_like(value_local) 
-
+                            count += torch.ones_like(value_local)
                     if use_dilated_sampling:
                         views = [[h, w] for h in range(current_scale_num) for w in range(current_scale_num)]
                         views_batch = [views[i : i + view_batch_size] for i in range(0, len(views), view_batch_size)]
-                        
                         h_pad = (current_scale_num - (latents.size(2) % current_scale_num)) % current_scale_num
                         w_pad = (current_scale_num - (latents.size(3) % current_scale_num)) % current_scale_num
                         latents_ = F.pad(latents, (w_pad, 0, h_pad, 0), 'constant', 0)
-                        
                         count_global = torch.zeros_like(latents_)
                         value_global = torch.zeros_like(latents_)
-
                         if use_guassian:
                             c3 = 0.99 * cosine_factor ** cosine_scale_3 + 1e-2
                             std_, mean_ = latents_.std(), latents_.mean()
@@ -1127,47 +983,31 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                             latents_gaussian = (latents_gaussian - latents_gaussian.mean()) / latents_gaussian.std() * std_ + mean_
                         else:
                             latents_gaussian = latents_
-
                         for j, batch_view in enumerate(views_batch):
-                            
                             latents_for_view = torch.cat(
-                                [
-                                    latents_[:, :, h::current_scale_num, w::current_scale_num]
-                                    for h, w in batch_view
-                                ]
+                                [latents_[:, :, h::current_scale_num, w::current_scale_num] for h, w in batch_view]
                             )
-
                             latents_for_view_gaussian = torch.cat(
-                                [
-                                    latents_gaussian[:, :, h::current_scale_num, w::current_scale_num]
-                                    for h, w in batch_view
-                                ]
+                                [latents_gaussian[:, :, h::current_scale_num, w::current_scale_num] for h, w in batch_view]
                             )
-
                             if shuffle:
                                 shape = latents_for_view.shape
-                                # 수정: range(...)의 괄호를 추가합니다.
+                                # 수정: range(...) 괄호 추가
                                 shuffle_index = torch.stack([torch.randperm(shape[0]) for _ in range(latents_for_view.reshape(-1).shape[0]//shape[0])])
                                 shuffle_index = shuffle_index.view(shape[1], shape[2], shape[3], shape[0])
                                 original_index = torch.zeros_like(shuffle_index).scatter_(3, shuffle_index, torch.arange(shape[0]).repeat(shape[1], shape[2], shape[3], 1))
                                 shuffle_index = shuffle_index.permute(3, 0, 1, 2).to(device)
                                 original_index = original_index.permute(3, 0, 1, 2).to(device)
                                 latents_for_view_gaussian = latents_for_view_gaussian.gather(0, shuffle_index)
-                        
                             vb_size = latents_for_view.size(0)
-
                             latent_model_input = latents_for_view_gaussian
-                            latent_model_input = (
-                                latent_model_input.repeat_interleave(2, dim=0)
-                                if do_classifier_free_guidance
-                                else latent_model_input
-                            )
+                            latent_model_input = (latent_model_input.repeat_interleave(2, dim=0)
+                                                    if do_classifier_free_guidance
+                                                    else latent_model_input)
                             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
                             prompt_embeds_input = torch.cat([prompt_embeds] * vb_size)
                             add_text_embeds_input = torch.cat([add_text_embeds] * vb_size)
                             add_time_ids_input = torch.cat([add_time_ids] * vb_size)
-
                             added_cond_kwargs = {"text_embeds": add_text_embeds_input, "time_ids": add_time_ids_input}
                             noise_pred = self.unet(
                                 latent_model_input,
@@ -1176,28 +1016,19 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                                 added_cond_kwargs=added_cond_kwargs,
                                 return_dict=False,
                             )[0]
-
                             if do_classifier_free_guidance:
                                 noise_pred_uncond, noise_pred_text = noise_pred[::2], noise_pred[1::2]
                                 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
                             if do_classifier_free_guidance and guidance_rescale > 0.0:
                                 noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
                             if shuffle:
                                 noise_pred = noise_pred.gather(0, original_index)
-
                             self.scheduler._init_step_index(t)
                             latents_denoised_batch = self.scheduler.step(noise_pred, t, latents_for_view, **extra_step_kwargs, return_dict=False)[0]
-
-                            for latents_view_denoised, (h, w) in zip(
-                                latents_denoised_batch.chunk(vb_size), batch_view
-                            ):
+                            for latents_view_denoised, (h, w) in zip(latents_denoised_batch.chunk(vb_size), batch_view):
                                 value_global[:, :, h::current_scale_num, w::current_scale_num] += latents_view_denoised
                                 count_global[:, :, h::current_scale_num, w::current_scale_num] += 1
-                                
                         value_global = value_global[:, :, h_pad:, w_pad:]
-
                         if use_multidiffusion:
                             c2 = cosine_factor ** cosine_scale_2
                             value += value_global * c2
@@ -1205,15 +1036,12 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                         else:
                             value += value_global
                             count += torch.ones_like(value_global)
-                                
                     latents = torch.where(count > 0, value / count, value)
-
                     if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                         progress_bar.update()
                         if callback is not None and i % callback_steps == 0:
                             step_idx = i // getattr(self.scheduler, "order", 1)
                             callback(step_idx, t, latents)
-
                 latents = (latents - latents.mean()) / latents.std() * anchor_std + anchor_mean
                 if self.lowvram:
                     latents = latents.cpu()
@@ -1224,29 +1052,23 @@ class AccDiffusionSDXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoade
                         needs_upcasting = False
                         self.unet.cpu()
                         self.vae.to(device)
-                    
                     if needs_upcasting:
                         self.upcast_vae()
                         latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
-
                     print("### Phase {} Decoding ###".format(current_scale_num))
                     if current_height > 2048 or current_width > 2048:
                         self.enable_vae_tiling()
                         image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
                     else:
                         image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-                    
                     image = self.image_processor.postprocess(image, output_type=output_type)
                     image[0].save(f'{result_path}/AccDiffusion_{current_scale_num}.png')
                     output_images.append(image[0])
-
                     if needs_upcasting:
                         self.vae.to(dtype=torch.float16)
                 else:
                     image = latents
-        
         self.maybe_free_model_hooks()
-
         return output_images
 
 
@@ -1280,9 +1102,7 @@ if __name__ == "__main__":
     ## others ##
     parser.add_argument('--debug', default=False, action='store_true')
     parser.add_argument('--experiment_name', default="AccDiffusion")
-
     args = parser.parse_args()
-
     pipe = AccDiffusionSDXLPipeline.from_pretrained(args.model_ckpt, torch_dtype=torch.float16).to("cuda")
     
     @spaces.GPU(duration=200)
@@ -1294,11 +1114,9 @@ if __name__ == "__main__":
                                   "n_cross_replace": {"default_": 1.0, "confetti": 0.8},
                                   }
         generator = torch.Generator(device='cuda').manual_seed(seed)
-        
         print(f"Prompt: {prompt}")
         md5_hash = hashlib.md5(prompt.encode()).hexdigest()
         result_path = f"./output/{args.experiment_name}/{md5_hash}/{width}_{height}_{seed}/"
-
         images = pipe(
             prompt, 
             negative_prompt=args.negative_prompt, 
@@ -1326,11 +1144,9 @@ if __name__ == "__main__":
             debug=args.debug, save_attention_map=args.save_attention_map, use_md_prompt=use_md_prompt, c=args.c
         )
         print(images)
-
         return images
     
     MAX_SEED = np.iinfo(np.int32).max
-
     css = """
     body {
         background: linear-gradient(135deg, #2c3e50, #4ca1af);
@@ -1353,7 +1169,6 @@ if __name__ == "__main__":
         visibility: hidden;
     }
     """
-
     with gr.Blocks(css=css) as demo:
         with gr.Column(elem_id="col-container"):
             gr.Markdown("<h1>AccDiffusion: Advanced AI Art Generator</h1>")
@@ -1363,7 +1178,6 @@ if __name__ == "__main__":
             with gr.Row():
                 prompt = gr.Textbox(label="Prompt", placeholder="예: A surreal landscape with floating islands and vibrant colors.", lines=2, scale=4)
                 submit_btn = gr.Button("Generate", scale=1)
-            
             with gr.Accordion("Advanced Settings", open=False):
                 with gr.Row():
                     resolution = gr.Radio(
@@ -1386,8 +1200,7 @@ if __name__ == "__main__":
                     use_progressive_upscaling = gr.Checkbox(label="Use Progressive Upscaling", value=False)
                     shuffle = gr.Checkbox(label="Shuffle", value=False)
                     use_md_prompt = gr.Checkbox(label="Use MD Prompt", value=False)
-            
-            output_images = gr.Gallery(label="Output Images", format="png").style(grid=[2], height="auto")
+            output_images = gr.Gallery(label="Output Images", format="png")
             gr.Markdown("### Example Prompts")
             gr.Examples(
                 examples = [