[hf staff] UI suggestion - stream the denoising process

app.py

@@ -36,6 +36,72 @@ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 MODEL_ID = os.environ.get("MODEL_ID", "Photoroom/prxpixel-t2i")
 HF_TOKEN = os.environ.get("HF_TOKEN")
 MAX_SEED = 2**31 - 1
+PREVIEW_SIZE = 512  # downscaled size for the per-step preview stream
+
+
+class StreamingPRXPixelPipeline(PRXPixelPipeline):
+    """PRXPixel pipeline whose denoising loop is a generator that yields the per-step x0 estimate.
+
+    Identical math to the parent `__call__` (same prompt encoding, noise_scale init, flow-matching
+    velocity conversion and scheduler step) — it just `yield`s the model's current clean-image
+    prediction x0 each step. Because PRXPixel is pixel-space, x0 is already an RGB image in [-1, 1],
+    so we can show it directly with no VAE decode. The final image is the running sample, matching
+    what the standard `__call__` returns.
+    """
+
+    @torch.inference_mode()
+    def stream(self, prompt, negative_prompt="", height=1024, width=1024,
+               num_inference_steps=28, guidance_scale=1.0, generator=None):
+        device = self._execution_device
+        self._guidance_scale = guidance_scale
+
+        text_embeddings, cross_attn_mask, uncond_text_embeddings, uncond_cross_attn_mask = self.encode_prompt(
+            prompt, device,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt or "",
+            num_images_per_prompt=1,
+        )
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        latents = self.prepare_latents(
+            1, self.transformer.config.in_channels, height, width,
+            text_embeddings.dtype, device, generator, None,
+        )
+        if self.do_classifier_free_guidance:
+            ca_embed = torch.cat([uncond_text_embeddings, text_embeddings], dim=0)
+            ca_mask = (
+                torch.cat([uncond_cross_attn_mask, cross_attn_mask], dim=0)
+                if cross_attn_mask is not None and uncond_cross_attn_mask is not None else None
+            )
+        else:
+            ca_embed, ca_mask = text_embeddings, cross_attn_mask
+
+        total = len(timesteps)
+        for i, t in enumerate(timesteps):
+            if self.do_classifier_free_guidance:
+                latents_in = torch.cat([latents, latents], dim=0)
+                t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).repeat(2).to(device)
+            else:
+                latents_in = latents
+                t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).to(device)
+
+            pred = self.transformer(
+                hidden_states=latents_in, timestep=t_cont,
+                encoder_hidden_states=ca_embed, attention_mask=ca_mask, return_dict=False,
+            )[0]
+            if self.do_classifier_free_guidance:
+                pred_uncond, pred_text = pred.chunk(2, dim=0)
+                pred = pred_uncond + guidance_scale * (pred_text - pred_uncond)
+
+            # `pred` is the predicted clean image x0 in [-1, 1] — the model's current guess.
+            yield i + 1, total, self.image_processor.postprocess(pred.float(), output_type="pil")[0]
+
+            t_x = torch.clamp(t.float() / self.scheduler.config.num_train_timesteps, min=0.05)
+            latents = self.scheduler.step((latents - pred) / t_x, t, latents).prev_sample
+
+        # Final image is the running sample, matching the standard pipeline output.
+        yield total, total, self.image_processor.postprocess(latents.float(), output_type="pil")[0]
+
 
 # --- Safety guard ---------------------------------------------------------
 # Small open-source NSFW prompt classifier (safe/nsfw) + a keyword blocklist for violence/gore
@@ -49,7 +115,7 @@ VIOLENCE_WORDS = [
 ]
 
 print(f"Loading {MODEL_ID} on CPU ...")
-pipe = PRXPixelPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN)
+pipe = StreamingPRXPixelPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN)
 print(f"Loading safety classifier {SAFETY_MODEL} on CPU ...")
 safety_tokenizer = AutoTokenizer.from_pretrained(SAFETY_MODEL)
 safety_model = AutoModelForSequenceClassification.from_pretrained(SAFETY_MODEL)
@@ -84,9 +150,8 @@ def generate(prompt, negative_prompt, steps, guidance_scale, shift, resolution,
     pipe.scheduler = FlowMatchEulerDiscreteScheduler(shift=float(shift))
 
     res = int(resolution)
-    # The pipeline already denormalizes to a proper PIL image. (Do NOT post-process the latents
-    # with (x+1)/2 yourself — output_type="pt" is already in [0,1], so re-normalizing washes it out.)
-    image = pipe(
+    # Stream the per-step x0 (clean-image) estimate so you can watch the image form, then the final.
+    for i, total, x0 in pipe.stream(
         prompt,
         negative_prompt=negative_prompt or "",
         height=res,
@@ -94,8 +159,13 @@ def generate(prompt, negative_prompt, steps, guidance_scale, shift, resolution,
         num_inference_steps=int(steps),
         guidance_scale=float(guidance_scale),
         generator=torch.Generator("cuda").manual_seed(seed),
-    ).images[0]
-    return image, seed
+    ):
+        if i < total:
+            preview = x0.copy()
+            preview.thumbnail((PREVIEW_SIZE, PREVIEW_SIZE))
+            yield preview, seed, f"Step {i} / {total}"
+        else:
+            yield x0, seed, f"Done — {total} steps"  # final, full resolution
 
 
 # (caption, prompt) for each pre-rendered example image examples/ex{i}.png
@@ -146,6 +216,9 @@ with gr.Blocks(title="PRX Pixel") as demo:
                 randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
             run = gr.Button("Generate", variant="primary")
         with gr.Column():
+            # The output now updates live with the model's per-step x0 estimate (pixel space → no decode),
+            # so you can watch the image form, blurry → sharp. Final frame is full resolution.
+            step_status = gr.Markdown("")
             out_image = gr.Image(label="Output", type="pil")
             out_seed = gr.Number(label="Seed used", interactive=False)
 
@@ -168,7 +241,7 @@ with gr.Blocks(title="PRX Pixel") as demo:
     run.click(
         generate,
         inputs=[prompt, negative_prompt, steps, guidance_scale, shift, resolution, seed, randomize_seed],
-        outputs=[out_image, out_seed],
+        outputs=[out_image, out_seed, step_status],
     )