Update app_quant_latent.py

app_quant_latent.py

@@ -253,127 +253,109 @@ import io
 logs = []
 latent_gallery = []
 def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
+image_seq_len,
+base_seq_len: int = 256,
+max_seq_len: int = 4096,
+base_shift: float = 0.5,
+max_shift: float = 1.15,
 ):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+def retrieve_timesteps(
+scheduler,
+num_inference_steps: int = None,
+device: str = None,
+timesteps: list = None,
+sigmas: list = None,
+**kwargs,
+):
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of timesteps or sigmas can be passed")
+ if timesteps is not None:
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
 
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed):
 
+ 
+ generator = torch.Generator(device).manual_seed(int(seed))
 
- try:
-    device = pipe._execution_device
-    generator = torch.Generator(device).manual_seed(int(seed))
+# Encode prompt
+ prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(prompt)
 
-    # 1. Encode prompt
-    prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(
-        prompt=prompt,
-        negative_prompt=None,
-        do_classifier_free_guidance=True,
-        device=device,
-    )
+ batch_size = len(prompt_embeds)
+ num_images_per_prompt = 1
+ actual_batch_size = batch_size * num_images_per_prompt
+ num_channels_latents = pipe.transformer.in_channels
 
-    batch_size = 1
-    num_images_per_prompt = 1
-    actual_batch_size = batch_size * num_images_per_prompt
-    num_channels_latents = pipe.transformer.in_channels
-
-    # 2. Prepare latents
-    latents = pipe.prepare_latents(
-        actual_batch_size,
-        num_channels_latents,
-        height,
-        width,
-        torch.float32,
-        device,
-        generator,
-        latents=None,
-    )
+ # Prepare latents
+ latents = pipe.prepare_latents(
+    actual_batch_size, num_channels_latents, height, width, torch.float32, device, generator
+ )
 
-    # Repeat prompt embeddings for multiple images
+ # Repeat embeddings for multiple images per prompt
+ if num_images_per_prompt > 1:
     prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)]
-    negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
-
-    # 3. Prepare timesteps
-    image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
-    mu = calculate_shift(
-        image_seq_len,
-        pipe.scheduler.config.get("base_image_seq_len", 256),
-        pipe.scheduler.config.get("max_image_seq_len", 4096),
-        pipe.scheduler.config.get("base_shift", 0.5),
-        pipe.scheduler.config.get("max_shift", 1.15),
-    )
-    pipe.scheduler.sigma_min = 0.0
-    timesteps, num_inference_steps = retrieve_timesteps(pipe.scheduler, steps, device, sigmas=None, mu=mu)
-
-    # 4. Denoising loop
-    with pipe.progress_bar(total=num_inference_steps) as progress_bar:
-        for i, t in enumerate(timesteps):
-            timestep = t.expand(latents.shape[0])
-            timestep = (1000 - timestep) / 1000
-
-            # CFG
-            latents_typed = latents.to(pipe.transformer.dtype)
-            latent_model_input = latents_typed.repeat(2, 1, 1, 1).unsqueeze(2)
-            prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds
-            timestep_model_input = timestep.repeat(2)
-
-            latent_model_input_list = list(latent_model_input.unbind(dim=0))
-            model_out_list = pipe.transformer(
-                latent_model_input_list, timestep_model_input, prompt_embeds_model_input, return_dict=False
-            )[0]
-
-            # Perform CFG
-            pos_out = model_out_list[:actual_batch_size]
-            neg_out = model_out_list[actual_batch_size:]
-            noise_pred = torch.stack([p + pipe.guidance_scale * (p - n) for p, n in zip(pos_out, neg_out)], dim=0)
-
-            noise_pred = noise_pred.squeeze(2)
-            noise_pred = -noise_pred
-            latents = pipe.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0]
-
-            # Store each latent step for gallery
-            latent_gallery.append(latents.clone().detach().cpu())
-            progress_bar.update()
-
-    # 5. Decode final image
-    latents_dec = latents.to(pipe.vae.dtype)
-    latents_dec = (latents_dec / pipe.vae.config.scaling_factor) + getattr(pipe.vae.config, "shift_factor", 0.0)
-
-    # Squeeze extra dim if present
-    if latents_dec.dim() == 5 and latents_dec.shape[2] == 1:
-        latents_dec = latents_dec.squeeze(2)
-
-    image = pipe.vae.decode(latents_dec, return_dict=False)[0]
-    final_image = pipe.image_processor.postprocess(image, output_type="pil")
-
-    # Decode latent gallery steps to images (optional)
-    gallery_images = []
-    for idx, lat in enumerate(latent_gallery):
-        try:
-            lat = lat.to(pipe.vae.dtype)
-            if lat.dim() == 5 and lat.shape[2] == 1:
-                lat = lat.squeeze(2)
-            img = pipe.vae.decode(lat, return_dict=False)[0]
-            img = pipe.image_processor.postprocess(img, output_type="pil")
-            gallery_images.append(img[0])
-        except Exception as e:
-            logs.append(f"⚠️ Failed to decode latent step {idx}: {e}")
-
-    return final_image[0], gallery_images, "\n".join(logs)
-
- except Exception as e:
-    return None, [], f"❌ Inference error: {e}"
-
-
+    if pipe.do_classifier_free_guidance and negative_prompt_embeds:
+        negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
+
+ image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
+ mu = calculate_shift(image_seq_len)
+
+ pipe.scheduler.sigma_min = 0.0
+ scheduler_kwargs = {"mu": mu}
+ timesteps, num_inference_steps = retrieve_timesteps(pipe.scheduler, steps, device, **scheduler_kwargs)
+
+ # Denoising loop
+ for i, t in enumerate(timesteps):
+    timestep = t.expand(latents.shape[0])
+    timestep = (1000 - timestep) / 1000
+    t_norm = timestep[0].item()
+    apply_cfg = pipe.do_classifier_free_guidance and pipe.guidance_scale > 0
+
+    if apply_cfg:
+        latent_model_input = latents.to(pipe.transformer.dtype).repeat(2, 1, 1, 1).unsqueeze(2)
+        prompt_input = prompt_embeds + negative_prompt_embeds
+        timestep_input = timestep.repeat(2)
+    else:
+        latent_model_input = latents.to(pipe.transformer.dtype).unsqueeze(2)
+        prompt_input = prompt_embeds
+        timestep_input = timestep
+
+    latent_list = list(latent_model_input.unbind(0))
+    model_out_list = pipe.transformer(latent_list, timestep_input, prompt_input, return_dict=False)[0]
+
+    if apply_cfg:
+        pos_out = model_out_list[:actual_batch_size]
+        neg_out = model_out_list[actual_batch_size:]
+        noise_pred = torch.stack([p + pipe.guidance_scale * (p - n) for p, n in zip(pos_out, neg_out)])
+    else:
+        noise_pred = torch.stack([t.float() for t in model_out_list], 0)
+
+    noise_pred = noise_pred.squeeze(2)
+    noise_pred = -noise_pred
+    latents = pipe.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0]
+
+ # Decode final image
+ latents = latents.to(pipe.vae.dtype)
+ latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
+ image = pipe.vae.decode(latents, return_dict=False)[0]
+ image = pipe.image_processor.postprocess(image, output_type="pil")
+
+ return image, None, None
 
 
 # ============================================================