Unverified Commit ba6008ab authored by Linoy Tsaban's avatar Linoy Tsaban Committed by GitHub
Browse files

[HiDream] code example (#11317)

parent a8f5134c
...@@ -36,7 +36,7 @@ EXAMPLE_DOC_STRING = """ ...@@ -36,7 +36,7 @@ EXAMPLE_DOC_STRING = """
```py ```py
>>> import torch >>> import torch
>>> from transformers import PreTrainedTokenizerFast, LlamaForCausalLM >>> from transformers import PreTrainedTokenizerFast, LlamaForCausalLM
>>> from diffusers import UniPCMultistepScheduler, HiDreamImagePipeline, HiDreamImageTransformer2DModel >>> from diffusers import UniPCMultistepScheduler, HiDreamImagePipeline
>>> scheduler = UniPCMultistepScheduler( >>> scheduler = UniPCMultistepScheduler(
... flow_shift=3.0, prediction_type="flow_prediction", use_flow_sigmas=True ... flow_shift=3.0, prediction_type="flow_prediction", use_flow_sigmas=True
...@@ -50,16 +50,11 @@ EXAMPLE_DOC_STRING = """ ...@@ -50,16 +50,11 @@ EXAMPLE_DOC_STRING = """
... torch_dtype=torch.bfloat16, ... torch_dtype=torch.bfloat16,
... ) ... )
>>> transformer = HiDreamImageTransformer2DModel.from_pretrained(
... "HiDream-ai/HiDream-I1-Full", subfolder="transformer", torch_dtype=torch.bfloat16
... )
>>> pipe = HiDreamImagePipeline.from_pretrained( >>> pipe = HiDreamImagePipeline.from_pretrained(
... "HiDream-ai/HiDream-I1-Full", ... "HiDream-ai/HiDream-I1-Full",
... scheduler=scheduler, ... scheduler=scheduler,
... tokenizer_4=tokenizer_4, ... tokenizer_4=tokenizer_4,
... text_encoder_4=text_encoder_4, ... text_encoder_4=text_encoder_4,
... transformer=transformer,
... torch_dtype=torch.bfloat16, ... torch_dtype=torch.bfloat16,
... ) ... )
>>> pipe.enable_model_cpu_offload() >>> pipe.enable_model_cpu_offload()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment