Spaces:
Paused
Paused
Update app_quant_latent.py
Browse files- app_quant_latent.py +57 -45
app_quant_latent.py
CHANGED
|
@@ -579,7 +579,6 @@ def upload_latents_to_hf(latent_dict, filename="latents.pt"):
|
|
| 579 |
os.remove(local_path)
|
| 580 |
raise e
|
| 581 |
|
| 582 |
-
|
| 583 |
@spaces.GPU
|
| 584 |
def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
|
| 585 |
LOGS = []
|
|
@@ -590,46 +589,36 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
|
|
| 590 |
latent_gallery = []
|
| 591 |
final_gallery = []
|
| 592 |
|
| 593 |
-
# ---
|
| 594 |
try:
|
| 595 |
latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
|
| 596 |
|
| 597 |
-
#
|
| 598 |
-
|
| 599 |
-
|
| 600 |
-
with torch.no_grad():
|
| 601 |
-
# 1️⃣ Try normal VAE decode if available
|
| 602 |
-
if hasattr(pipe, "vae") and hasattr(pipe.vae, "decode"):
|
| 603 |
-
try:
|
| 604 |
-
latent_img_tensor = pipe.vae.decode(latents).sample # [1,3,H,W]
|
| 605 |
-
latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
|
| 606 |
-
latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
|
| 607 |
-
latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
|
| 608 |
-
except Exception as e1:
|
| 609 |
-
LOGS.append(f"⚠️ VAE decode failed: {e1}")
|
| 610 |
-
|
| 611 |
-
# 2️⃣ Collapse first 3 channels if decode failed
|
| 612 |
-
if latent_img is placeholder and latents.shape[1] >= 3:
|
| 613 |
-
ch = latents[0, :3, :, :]
|
| 614 |
-
ch = (ch - ch.min()) / (ch.max() - ch.min() + 1e-8)
|
| 615 |
-
latent_img = Image.fromarray((ch.permute(1, 2, 0).cpu().numpy() * 255).astype('uint8'))
|
| 616 |
-
|
| 617 |
-
# 3️⃣ Collapse all channels to mean -> replicate to RGB
|
| 618 |
-
if latent_img is placeholder:
|
| 619 |
-
mean_ch = latents[0].mean(dim=0, keepdim=True) # [1,H,W]
|
| 620 |
-
mean_ch = (mean_ch - mean_ch.min()) / (mean_ch.max() - mean_ch.min() + 1e-8)
|
| 621 |
-
latent_img = Image.fromarray(
|
| 622 |
-
torch.cat([mean_ch]*3, dim=0).permute(1,2,0).cpu().numpy().astype('uint8')
|
| 623 |
-
)
|
| 624 |
-
|
| 625 |
-
except Exception as e:
|
| 626 |
-
LOGS.append(f"⚠️ Latent to image conversion failed: {e}")
|
| 627 |
-
latent_img = placeholder
|
| 628 |
|
| 629 |
-
|
| 630 |
-
|
|
|
|
| 631 |
|
| 632 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 633 |
latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
|
| 634 |
try:
|
| 635 |
hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
|
|
@@ -642,7 +631,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
|
|
| 642 |
latent_gallery.append(placeholder)
|
| 643 |
yield None, latent_gallery, LOGS
|
| 644 |
|
| 645 |
-
# --- Final image:
|
| 646 |
try:
|
| 647 |
output = pipe(
|
| 648 |
prompt=prompt,
|
|
@@ -663,9 +652,9 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
|
|
| 663 |
final_gallery.append(placeholder)
|
| 664 |
latent_gallery.append(placeholder)
|
| 665 |
yield placeholder, latent_gallery, LOGS
|
| 666 |
-
# this
|
| 667 |
@spaces.GPU
|
| 668 |
-
def
|
| 669 |
LOGS = []
|
| 670 |
device = "cuda"
|
| 671 |
generator = torch.Generator(device).manual_seed(int(seed))
|
|
@@ -678,14 +667,36 @@ def generate_image_workswell(prompt, height, width, steps, seed, guidance_scale=
|
|
| 678 |
try:
|
| 679 |
latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
|
| 680 |
|
| 681 |
-
# Decode latent tensor to PIL for preview
|
|
|
|
| 682 |
try:
|
| 683 |
with torch.no_grad():
|
| 684 |
-
|
| 685 |
-
|
| 686 |
-
|
| 687 |
-
|
| 688 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 689 |
latent_img = placeholder
|
| 690 |
|
| 691 |
latent_gallery.append(latent_img)
|
|
@@ -725,6 +736,7 @@ def generate_image_workswell(prompt, height, width, steps, seed, guidance_scale=
|
|
| 725 |
final_gallery.append(placeholder)
|
| 726 |
latent_gallery.append(placeholder)
|
| 727 |
yield placeholder, latent_gallery, LOGS
|
|
|
|
| 728 |
|
| 729 |
|
| 730 |
|
|
|
|
| 579 |
os.remove(local_path)
|
| 580 |
raise e
|
| 581 |
|
|
|
|
| 582 |
@spaces.GPU
|
| 583 |
def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
|
| 584 |
LOGS = []
|
|
|
|
| 589 |
latent_gallery = []
|
| 590 |
final_gallery = []
|
| 591 |
|
| 592 |
+
# --- Generate latent previews in a loop ---
|
| 593 |
try:
|
| 594 |
latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
|
| 595 |
|
| 596 |
+
# Convert latents to float32 if necessary
|
| 597 |
+
if latents.dtype != torch.float32:
|
| 598 |
+
latents = latents.float()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 599 |
|
| 600 |
+
# Loop for multiple previews before final image
|
| 601 |
+
num_previews = min(10, steps) # show ~10 previews
|
| 602 |
+
preview_steps = torch.linspace(0, 1, num_previews)
|
| 603 |
|
| 604 |
+
for i, alpha in enumerate(preview_steps):
|
| 605 |
+
try:
|
| 606 |
+
with torch.no_grad():
|
| 607 |
+
# Simple noise interpolation for preview (simulate denoising progress)
|
| 608 |
+
preview_latent = latents * alpha + torch.randn_like(latents) * (1 - alpha)
|
| 609 |
+
# Decode to PIL
|
| 610 |
+
latent_img_tensor = pipe.vae.decode(preview_latent).sample # [1,3,H,W]
|
| 611 |
+
latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
|
| 612 |
+
latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
|
| 613 |
+
latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
|
| 614 |
+
except Exception as e:
|
| 615 |
+
LOGS.append(f"⚠️ Latent preview decode failed: {e}")
|
| 616 |
+
latent_img = placeholder
|
| 617 |
+
|
| 618 |
+
latent_gallery.append(latent_img)
|
| 619 |
+
yield None, latent_gallery, LOGS # update Gradio with intermediate preview
|
| 620 |
+
|
| 621 |
+
# Save final latents to HF
|
| 622 |
latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
|
| 623 |
try:
|
| 624 |
hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
|
|
|
|
| 631 |
latent_gallery.append(placeholder)
|
| 632 |
yield None, latent_gallery, LOGS
|
| 633 |
|
| 634 |
+
# --- Final image: untouched standard pipeline ---
|
| 635 |
try:
|
| 636 |
output = pipe(
|
| 637 |
prompt=prompt,
|
|
|
|
| 652 |
final_gallery.append(placeholder)
|
| 653 |
latent_gallery.append(placeholder)
|
| 654 |
yield placeholder, latent_gallery, LOGS
|
| 655 |
+
# this is astable vesopn tha can gen final and a noise to latent
|
| 656 |
@spaces.GPU
|
| 657 |
+
def generate_image0(prompt, height, width, steps, seed, guidance_scale=0.0):
|
| 658 |
LOGS = []
|
| 659 |
device = "cuda"
|
| 660 |
generator = torch.Generator(device).manual_seed(int(seed))
|
|
|
|
| 667 |
try:
|
| 668 |
latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
|
| 669 |
|
| 670 |
+
# Decode latent tensor to PIL for preview with robust fallbacks
|
| 671 |
+
latent_img = placeholder
|
| 672 |
try:
|
| 673 |
with torch.no_grad():
|
| 674 |
+
# 1️⃣ Try normal VAE decode if available
|
| 675 |
+
if hasattr(pipe, "vae") and hasattr(pipe.vae, "decode"):
|
| 676 |
+
try:
|
| 677 |
+
latent_img_tensor = pipe.vae.decode(latents).sample # [1,3,H,W]
|
| 678 |
+
latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
|
| 679 |
+
latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
|
| 680 |
+
latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
|
| 681 |
+
except Exception as e1:
|
| 682 |
+
LOGS.append(f"⚠️ VAE decode failed: {e1}")
|
| 683 |
+
|
| 684 |
+
# 2️⃣ Collapse first 3 channels if decode failed
|
| 685 |
+
if latent_img is placeholder and latents.shape[1] >= 3:
|
| 686 |
+
ch = latents[0, :3, :, :]
|
| 687 |
+
ch = (ch - ch.min()) / (ch.max() - ch.min() + 1e-8)
|
| 688 |
+
latent_img = Image.fromarray((ch.permute(1, 2, 0).cpu().numpy() * 255).astype('uint8'))
|
| 689 |
+
|
| 690 |
+
# 3️⃣ Collapse all channels to mean -> replicate to RGB
|
| 691 |
+
if latent_img is placeholder:
|
| 692 |
+
mean_ch = latents[0].mean(dim=0, keepdim=True) # [1,H,W]
|
| 693 |
+
mean_ch = (mean_ch - mean_ch.min()) / (mean_ch.max() - mean_ch.min() + 1e-8)
|
| 694 |
+
latent_img = Image.fromarray(
|
| 695 |
+
torch.cat([mean_ch]*3, dim=0).permute(1,2,0).cpu().numpy().astype('uint8')
|
| 696 |
+
)
|
| 697 |
+
|
| 698 |
+
except Exception as e:
|
| 699 |
+
LOGS.append(f"⚠️ Latent to image conversion failed: {e}")
|
| 700 |
latent_img = placeholder
|
| 701 |
|
| 702 |
latent_gallery.append(latent_img)
|
|
|
|
| 736 |
final_gallery.append(placeholder)
|
| 737 |
latent_gallery.append(placeholder)
|
| 738 |
yield placeholder, latent_gallery, LOGS
|
| 739 |
+
# this version generate well for final and gives a tensor back for latent
|
| 740 |
|
| 741 |
|
| 742 |
|