Update app_quant_latent.py

app_quant_latent.py

@@ -554,24 +554,47 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 # --------------------------
 # Main generation function (kept exactly as your logic)
 # --------------------------
+from huggingface_hub import HfApi, HfFolder
+import torch
+import os
+
+HF_REPO_ID = "rahul7star/Zstudio-latent"  # Model repo
+HF_TOKEN = HfFolder.get_token()  # Make sure you are logged in via `huggingface-cli login`
+
+def upload_latents_to_hf(latent_dict, filename="latents.pt"):
+    local_path = f"/tmp/{filename}"
+    torch.save(latent_dict, local_path)
+    try:
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=filename,
+            repo_id=HF_REPO_ID,
+            token=HF_TOKEN,
+            repo_type="model"  # since this is a model repo
+        )
+        os.remove(local_path)
+        return f"https://huggingface.co/{HF_REPO_ID}/resolve/main/{filename}"
+    except Exception as e:
+        os.remove(local_path)
+        raise e
+
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
     device = "cuda"
     generator = torch.Generator(device).manual_seed(int(seed))
 
-    # placeholders
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []
     final_gallery = []
 
     try:
-        # --- Try advanced latent mode ---
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
+            # Step-wise denoising
             for i, t in enumerate(pipe.scheduler.timesteps):
-                # Step-wise denoising
                 with torch.no_grad():
                     noise_pred = pipe.unet(
                         latents,
@@ -580,28 +603,34 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
                     )["sample"]
                     latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"]
 
-                # Convert latent to preview image
-                try:
-                    latent_img = latent_to_image(latents, pipe.vae)[0]
-                except Exception:
-                    latent_img = placeholder
-
-                latent_gallery.append(latent_img)
-
-                # Yield intermediate update: latents updated, final gallery empty
-                yield None, latent_gallery, final_gallery, LOGS
+                # Convert latent to preview image every few steps
+                if i % max(1, len(pipe.scheduler.timesteps)//10) == 0 or i == len(pipe.scheduler.timesteps)-1:
+                    try:
+                        latent_img = latent_to_image(latents, pipe.vae)[0]
+                    except Exception:
+                        latent_img = placeholder
+                    latent_gallery.append(latent_img)
+                    yield None, latent_gallery, LOGS  # yield intermediate latents
 
             # Decode final image
             final_img = pipe.decode_latents(latents)[0]
             final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
-            yield final_img, latent_gallery, final_gallery, LOGS
+
+            # Save latents to dict and upload to HF
+            latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
+            try:
+                hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
+                LOGS.append(f"🔹 Latents uploaded: {hf_url}")
+            except Exception as e:
+                LOGS.append(f"⚠️ Failed to upload latents: {e}")
+
+            yield final_img, latent_gallery, LOGS
 
         except Exception as e:
             LOGS.append(f"⚠️ Advanced latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
-            # Standard pipeline fallback
             try:
                 output = pipe(
                     prompt=prompt,
@@ -613,22 +642,21 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
                 )
                 final_img = output.images[0]
                 final_gallery.append(final_img)
-                latent_gallery.append(final_img)  # optionally show in latent gallery as last step
+                latent_gallery.append(final_img)  # fallback latent preview
                 LOGS.append("✅ Standard pipeline succeeded.")
-                yield final_img, latent_gallery, final_gallery, LOGS
+                yield final_img, latent_gallery, LOGS
 
             except Exception as e2:
                 LOGS.append(f"❌ Standard pipeline failed: {e2}")
                 final_gallery.append(placeholder)
                 latent_gallery.append(placeholder)
-                yield placeholder, latent_gallery, final_gallery, LOGS
+                yield placeholder, latent_gallery, LOGS
 
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
         final_gallery.append(placeholder)
         latent_gallery.append(placeholder)
-        yield placeholder, latent_gallery, final_gallery, LOGS
-
+        yield placeholder, latent_gallery, LOGS