fp16 inference

README.md

@@ -26,6 +26,7 @@ pipeline = DiffusionPipeline.from_pretrained(
     vae=sd3_vae,
     custom_pipeline="StonyBrook-CVLab/PixCell-pipeline",
     trust_remote_code=True,
+    torch_dtype=torch.float16,
 )
 
 pipeline.to(device);
@@ -61,7 +62,8 @@ uni_model.to(device);
 ### Unconditional generation
 ```python
 uncond = pipeline.get_unconditional_embedding(1)
-samples = pipeline(uni_embeds=uncond, negative_uni_embeds=None, guidance_scale=1.0)
+with torch.amp.autocast('cuda'):
+    samples = pipeline(uni_embeds=uncond, negative_uni_embeds=None, guidance_scale=1.0)
 ```
 
 ### Conditional generation
@@ -87,5 +89,6 @@ print("Extracted UNI:", uni_emb.shape)
 # Get unconditional embedding for classifier-free guidance
 uncond = pipeline.get_unconditional_embedding(uni_emb.shape[0])
 # Generate new samples
-samples = pipeline(uni_embeds=uni_emb, negative_uni_embeds=uncond, guidance_scale=3., num_images_per_prompt=1).images
+with torch.amp.autocast('cuda'):
+    samples = pipeline(uni_embeds=uni_emb, negative_uni_embeds=uncond, guidance_scale=3., num_images_per_prompt=1).images
 ```