Update README.md

README.md

@@ -121,31 +121,85 @@ image.save("./alt.png")
 
 ![alt](https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/hub/alt.png)
 
-##  FlagAI Example
-
-以下示例将为文本输入`Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation` 在目录`./AltDiffusionOutputs`下生成图片结果。
-
-The following example will generate image results for text input `Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation` under the default output directory `./AltDiffusionOutputs`
+## Transformers Example
 
 ```python
+import os
 import torch
-from flagai.auto_model.auto_loader import AutoLoader
-from flagai.model.predictor.predictor import Predictor
-
-# Initialize 
-prompt = "Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-loader = AutoLoader(task_name="text2img", #contrastive learning
-                    model_name="AltDiffusion",
-                    model_dir="./checkpoints")
-
-model = loader.get_model()
-model.eval()
-model.to(device)
-predictor = Predictor(model)
-predictor.predict_generate_images(prompt)
+import transformers
+from transformers import BertPreTrainedModel
+from transformers.models.clip.modeling_clip import CLIPPreTrainedModel
+from transformers.models.xlm_roberta.tokenization_xlm_roberta import XLMRobertaTokenizer
+from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
+from diffusers import StableDiffusionPipeline
+from transformers import BertPreTrainedModel,BertModel,BertConfig
+import torch.nn as nn
+import torch
+from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig
+from transformers import XLMRobertaModel
+from transformers.activations import ACT2FN
+from typing import Optional
+
+
+class RobertaSeriesConfig(XLMRobertaConfig):
+    def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2,project_dim=768,pooler_fn='cls',learn_encoder=False, **kwargs):
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.project_dim = project_dim
+        self.pooler_fn = pooler_fn
+        # self.learn_encoder = learn_encoder
+
+class RobertaSeriesModelWithTransformation(BertPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+    base_model_prefix = 'roberta'
+    config_class= XLMRobertaConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = XLMRobertaModel(config)
+        self.transformation = nn.Linear(config.hidden_size, config.project_dim)
+        self.post_init()
+        
+    def get_text_embeds(self,bert_embeds,clip_embeds):
+        return self.merge_head(torch.cat((bert_embeds,clip_embeds)))
+
+    def set_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def forward(self, input_ids: Optional[torch.Tensor] = None) :
+        attention_mask = (input_ids != self.tokenizer.pad_token_id).to(torch.int64)
+        outputs = self.base_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+        
+        projection_state = self.transformation(outputs.last_hidden_state)
+        
+        return (projection_state,)
+
+model_path_encoder = "BAAI/RobertaSeriesModelWithTransformation"
+model_path_diffusion = "BAAI/AltDiffusion"
+device = "cuda"
+
+seed = 12345
+tokenizer = XLMRobertaTokenizer.from_pretrained(model_path_encoder, use_auth_token=True)
+tokenizer.model_max_length = 77
+
+text_encoder = RobertaSeriesModelWithTransformation.from_pretrained(model_path_encoder, use_auth_token=True)
+text_encoder.set_tokenizer(tokenizer)
+print("text encode loaded")
+pipe = StableDiffusionPipeline.from_pretrained(model_path_diffusion,
+                                               tokenizer=tokenizer,
+                                               text_encoder=text_encoder,
+                                               use_auth_token=True,
+                                               )
+print("diffusion pipeline loaded")
+pipe = pipe.to(device)
+
+prompt = "Thirty years old lee evans as a sad 19th century postman. detailed, soft focus, candle light, interesting lights, realistic, oil canvas, character concept art by munkácsy mihály, csók istván, john everett millais, henry meynell rheam, and da vinci"
+with torch.no_grad():
+    image = pipe(prompt, guidance_scale=7.5).images[0]  
+    
+image.save("3.png")
 ```