{"id":13488411,"url":"https://github.com/chaofengc/texforce","last_synced_at":"2025-08-21T15:33:36.145Z","repository":{"id":209460043,"uuid":"723888152","full_name":"chaofengc/TexForce","owner":"chaofengc","description":"Official PyTorch codes for \"Enhancing Diffusion Models with Text-Encoder Reinforcement Learning\", ECCV2024","archived":false,"fork":false,"pushed_at":"2024-08-13T13:21:16.000Z","size":66005,"stargazers_count":53,"open_issues_count":3,"forks_count":4,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-08T12:47:28.851Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/chaofengc.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2023-11-27T01:08:08.000Z","updated_at":"2025-03-26T04:44:38.000Z","dependencies_parsed_at":"2023-12-23T09:46:35.621Z","dependency_job_id":null,"html_url":"https://github.com/chaofengc/TexForce","commit_stats":null,"previous_names":["chaofengc/texforce"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/chaofengc/TexForce","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/chaofengc%2FTexForce","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/chaofengc%2FTexForce/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/chaofengc%2FTexForce/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/chaofengc%2FTexForce/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/chaofengc","download_url":"https://codeload.github.com/chaofengc/TexForce/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/chaofengc%2FTexForce/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":271500361,"owners_count":24770375,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-08-21T02:00:08.990Z","response_time":74,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-07-31T18:01:15.256Z","updated_at":"2025-08-21T15:33:31.131Z","avatar_url":"https://github.com/chaofengc.png","language":"Python","funding_links":[],"categories":["T2I Diffusion Model augmentation"],"sub_categories":[],"readme":"# Enhancing Diffusion Models with Text-Encoder Reinforcement Learning\n\nOfficial PyTorch codes for paper [Enhancing Diffusion Models with Text-Encoder Reinforcement Learning](https://arxiv.org/abs/2311.15657)\n\n[![arXiv](https://img.shields.io/badge/arXiv-Paper-\u003cCOLOR\u003e.svg)](https://arxiv.org/abs/2311.15657)\n\u003ca href=\"https://colab.research.google.com/drive/1pC6lzFR4hohwWyBgnkeZUghi2jqjl9gI?usp=sharing\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"google colab logo\"\u003e\u003c/a\u003e \n[![huggingface](https://img.shields.io/badge/HuggingFace-model-red.svg)](https://huggingface.co/chaofengc/sd-turbo_texforce)\n![visitors](https://visitor-badge.laobi.icu/badge?page_id=chaofengc/TexForce)\n\n![teaser_img](./assets/fig_teaser.jpg)\n\n## Requirements \u0026 Installation\n\n- Clone the repo and install required packages with \n```\n# git clone this repository\ngit clone https://github.com/chaofengc/TexForce.git\ncd TexForce \n\n# create new anaconda env\nconda create -n texforce python=3.8\nsource activate texforce \n\n# install python dependencies\npip3 install -r requirements.txt\n```\n\n## Results on SDXL-Turbo\n\nWe also applied our method to the recent model [sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo). The model is trained with [ImageReward](https://github.com/THUDM/ImageReward) feedback through direct back-propagation to save training time. Test with the following codes\n\n```\n## Note: sdturboxl requires latest diffusers installed from source with the following command\ngit clone https://github.com/huggingface/diffusers\ncd diffusers\npip install -e .\n```\n```\nfrom diffusers import AutoPipelineForText2Image\nimport torch\n\npipe = AutoPipelineForText2Image.from_pretrained(\"stabilityai/sdxl-turbo\", torch_dtype=torch.float16, variant=\"fp16\")\npipe = pipe.to(\"cuda\")\npipe.load_lora_weights('chaofengc/sdxl-turbo_texforce')\n\npt = ['a photo of a cat.']\nimg = pipe(prompt=pt, num_inference_steps=1, guidance_scale=0.0).images[0]\n```\n\nHere are some example results:\n\n\u003ctable\u003e\n\u003cthead\u003e\n  \u003ctr\u003e\n    \u003cth width=\"50%\"\u003esdxl-turbo\u003c/th\u003e\n    \u003cth width=\"50%\"\u003esdxl-turbo + TexForce\u003c/th\u003e\n  \u003c/tr\u003e\n\u003c/thead\u003e\n\u003ctbody\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      A photo of a cat.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n        \u003cimg src='assets/image_sdxlturbo_001.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      An astronaut riding a horse.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n        \u003cimg src='assets/image_sdxlturbo_002.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      water bottle.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n        \u003cimg src='assets/image_sdxlturbo_003.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/tbody\u003e\n\u003c/table\u003e\n\n\n## Results on SD-Turbo\n\nWe applied our method to the recent model [sdturbo](https://huggingface.co/stabilityai/sd-turbo). The model is trained with [Q-Instruct](https://github.com/Q-Future/Q-Instruct) feedback through direct back-propagation to save training time. Test with the following codes\n\n```\n## Note: sdturbo requires latest diffusers\u003e=0.24.0 with AutoPipelineForText2Image class\n\nfrom diffusers import AutoPipelineForText2Image\nfrom peft import PeftModel\nimport torch\n\npipe = AutoPipelineForText2Image.from_pretrained(\"stabilityai/sd-turbo\", torch_dtype=torch.float16, variant=\"fp16\")\npipe = pipe.to(\"cuda\")\nPeftModel.from_pretrained(pipe.text_encoder, 'chaofengc/sd-turbo_texforce')\n\npt = ['a photo of a cat.']\nimg = pipe(prompt=pt, num_inference_steps=1, guidance_scale=0.0).images[0]\n```\n\nHere are some example results:\n\n\u003ctable\u003e\n\u003cthead\u003e\n  \u003ctr\u003e\n    \u003cth width=\"50%\"\u003esd-turbo\u003c/th\u003e\n    \u003cth width=\"50%\"\u003esd-turbo + TexForce\u003c/th\u003e\n  \u003c/tr\u003e\n\u003c/thead\u003e\n\u003ctbody\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      A photo of a cat.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_001_sdturbo.jpg'\u003e\n    \u003c/td\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_001_sdturbo_texforce.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      A photo of a dog.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_002_sdturbo.jpg'\u003e\n    \u003c/td\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_002_sdturbo_texforce.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"2\"\u003e\n      A photo of a boy, colorful.\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_003_sdturbo.jpg'\u003e\n    \u003c/td\u003e\n    \u003ctd\u003e\n        \u003cimg src='assets/image_003_sdturbo_texforce.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n\n\u003c/tbody\u003e\n\u003c/table\u003e\n\n## Results on SD-1.4, SD-1.5, SD-2.1\n\nDue to code compatibility, you need to install the following diffusers first:\n```\npip uninstall diffusers\npip install diffusers==0.16.0\n```\n\nYou may simply load the pretrained lora weights with the following code block to improve performance of original stable diffusion model:\n```\nfrom diffusers import StableDiffusionPipeline\nfrom diffusers import DDIMScheduler \nfrom peft import PeftModel\nimport torch\n\ndef load_model_weights(pipe, weight_path, model_type):\n    if model_type == 'text+lora':\n        text_encoder = pipe.text_encoder\n        PeftModel.from_pretrained(text_encoder, weight_path)\n    elif model_type == 'unet+lora':\n        pipe.unet.load_attn_procs(weight_path)\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nmodel_id = \"CompVis/stable-diffusion-v1-4\"\npipe = StableDiffusionPipeline.from_pretrained(model_id, dtype=torch.float16)\npipe = pipe.to(\"cuda\")\npipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)\n\nload_model_weights(pipe, './lora_weights/sd14_refl/', 'unet+lora')\nload_model_weights(pipe, './lora_weights/sd14_texforce/', 'text+lora')\n\nprompt = ['a painting of a dog.']\nimg = pipe(prompt).images[0]\n\n```\n\nHere are some example results:\n\n\u003ctable\u003e\n\u003cthead\u003e\n  \u003ctr\u003e\n    \u003cth width=\"25%\"\u003eSDv1.4\u003c/th\u003e\n    \u003cth width=\"25%\"\u003eReFL\u003c/th\u003e\n    \u003cth width=\"25%\"\u003eTexForce\u003c/th\u003e\n    \u003cth width=\"25%\"\u003eReFL+TexForce\u003c/th\u003e\n  \u003c/tr\u003e\n\u003c/thead\u003e\n\u003ctbody\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n      astronaut drifting afloat in space, in the darkness away from anyone else, alone, black background dotted with stars, realistic\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n        \u003cimg src='assets/image_0021.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n      portrait of a cute cyberpunk cat, realistic, professional\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n        \u003cimg src='assets/image_0058.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n      a coffee mug made of cardboard\n    \u003c/td\u003e\n  \u003c/tr\u003e\n  \u003ctr\u003e\n    \u003ctd colspan=\"4\"\u003e\n        \u003cimg src='assets/image_0099.jpg'\u003e\n    \u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/tbody\u003e\n\u003c/table\u003e\n\n\n## Training\n\nWe rewrite the training codes based on [trl](https://github.com/huggingface/trl) with the latest diffusers library. \n\u003e [!NOTE]  \n\u003e The latest diffusers support simple loading of lora weights with `pipeline.load_lora_weights` after training.\n\nYou may train the model with the following command:\n\n### Example script for single prompt training\n```\naccelerate launch --num_processes 2 src/train_ddpo.py \\\n    --mixed_precision=\"fp16\" \\\n    --sample_num_steps 50 --train_timestep_fraction 0.5 \\\n    --num_epochs 40 \\\n    --sample_batch_size 4 --sample_num_batches_per_epoch 64 \\\n    --train_batch_size 4 --train_gradient_accumulation_steps 1 \\\n    --prompt=\"single\" --single_prompt_type=\"hand\" --reward_list=\"handdetreward\" \\\n    --per_prompt_stat_tracking=True \\\n    --tracker_project_name=\"texforce_hand\" \\\n    --log_with=\"tensorboard\"\n```\nThe supported prompts and reward functions are listed below:\n- prompts: `hand`, `face`, `color`, `count`, `comp`, `location`\n- rewards: `handdetreward`, `topiq_nr-face`, `imagereward` \n\n### Example script for complex multi-prompt training\n```\naccelerate launch --num_processes 2 src/train_ddpo.py \\\n    --mixed_precision=\"fp16\" \\\n    --sample_num_steps 50 --train_timestep_fraction 0.5 \\\n    --num_epochs 50 \\\n    --sample_batch_size 4 --sample_num_batches_per_epoch 128 \\\n    --train_batch_size 4 --train_gradient_accumulation_steps 4 \\\n    --prompt=\"imagereward\" --reward_list=\"imagereward\" \\\n    --per_prompt_stat_tracking=True \\\n    --tracker_project_name=\"texforce_imgreward\" \\\n    --log_with=\"tensorboard\"\n``` \nThe supported prompts and reward functions are:\n- prompts: `imagereward`, `hps`\n- rewards: `imagereward`, `hpsreward`, `laion_aes`\n\n## Citation\n\nIf you find this code useful for your research, please cite our paper:\n```\n@inproceedings{chen2024texforce,\n  title={Enhancing Diffusion Models with Text-Encoder Reinforcement Learning},\n  author={Chaofeng Chen and Annan Wang and Haoning Wu and Liang Liao and Wenxiu Sun and Qiong Yan and Weisi Lin},\n  year={2024},\n  booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},\n}\n```\n\n## License\n\nThis work is licensed under [NTU S-Lab License 1.0](./LICENCE_S-Lab) and a \u003ca rel=\"license\" href=\"http://creativecommons.org/licenses/by-nc-sa/4.0/\"\u003eCreative Commons Attribution-NonCommercial-ShareAlike 4.0 International License\u003c/a\u003e.\n\n\u003ca rel=\"license\" href=\"http://creativecommons.org/licenses/by-nc-sa/4.0/\"\u003e\u003cimg alt=\"Creative Commons License\" style=\"border-width:0\" src=\"https://i.creativecommons.org/l/by-nc-sa/4.0/88x31.png\" /\u003e\u003c/a\u003e\n\n\n## Acknowledgement\n\nThis project is largely based on [trl](https://github.com/huggingface/trl). The hand detection codes are taken from [Unified-Gesture-and-Fingertip-Detection](https://github.com/MahmudulAlam/Unified-Gesture-and-Fingertip-Detection). Many thanks to their great work :hugs:!\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fchaofengc%2Ftexforce","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fchaofengc%2Ftexforce","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fchaofengc%2Ftexforce/lists"}