[深度學習]利用Stable Diffusion輸入文字生成圖片

本文下方連結的文章，利用Stable Diffusion生成512 * 512大小的圖片。

輸入的文字是 dog flying in space，此模型需輸入英文句子才會準確生成。

參考文獻

連結該作者在Hugging Face公開的模型去做使用。

本文是在Colab上執行。Python版本及套件版本

Python version: 3.10.12
diffusers version: 0.29.2
transformers version: 4.42.4
scipy version: 1.14.0

程式範例

1.安裝套件

!pip install --upgrade diffusers transformers scipy &> /dev/null

2. 程式碼

短短幾行程式碼，即可載入Stable Diffusion模型來生成圖片。

只要修改prompt就可以改變要生成的圖片。

import torch
from diffusers import StableDiffusionPipeline

#設定模型 ID 和設備
model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

#載入模型
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)
#torch_dtype=torch.float16 ：單精度，較不佔記憶體

#生成圖像的描述
prompt = "dog flying in space"

#生成圖片
image = pipe(prompt).images[0]  ##預設圖片大小：512*768
#儲存圖片
image.save("dog_flying_in_space.png")

'''
高度和寬度都是8的倍數
低於 512 可能會導致影像品質降低, 超過 512 會出現重複影像區域（遺失全域一致性）
非正方形影像最佳方法：一個維度為 512px，另一維度大於512px。
'''

程式碼詳細說明

• 1. 導入相關套件

  import torch
  from diffusers import StableDiffusionPipeline

  • import torch：引入 PyTorch 庫，這是用於深度學習的熱門框架。
  • from diffusers import StableDiffusionPipeline：從 diffusers 庫中引入 StableDiffusionPipeline，這是用於圖像生成的管道。

• 2. 設定模型和設備

  model_id = "CompVis/stable-diffusion-v1-4"
  device = "cuda"

  • model_id：指定要使用的模型 ID，此處是 CompVis/stable-diffusion-v1-4。
  • device：此處設為 "cuda"，表示使用 GPU。

• 3. 載入模型

  # torch_dtype=torch.float16 ：單精度，較不佔記憶體
  pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
  pipe = pipe.to(device)

  • StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)：從預訓練模型載入管道，並設定 torch_dtype 為 torch.float16。使用 float16 可以減少記憶體佔用。
  • 這行程式碼會從 Hugging Face 下載並載入 CompVis/stable-diffusion-v1-4 模型，然後你可以使用這個模型進行圖像生成任務。
  • pipe.to(device)：指定設備（GPU）。

• 4. 生成圖像

  prompt = "dog flying in space"
  
  # 預設圖片大小：512*768
  '''
  高度和寬度都是8的倍數
  低於 512 可能會導致影像品質降低, 超過 512 會出現重複影像區域（遺失全域一致性）
  非正方形影像最佳方法：一個維度為 512px，另一維度大於512px。
  '''
  image = pipe(prompt).images[0]

  • prompt = "dog flying in space"：設置文本提示詞，這是生成圖像的描述。
  • image = pipe(prompt).images[0]：使用管道生成圖像。 pipe(prompt) 會返回一個結果對象， images[0] 表示取第一張生成的圖像。

• 5. 保存圖像

  image.save("astronaut_rides_horse.png")

  • image.save("astronaut_rides_horse.png")：將生成的圖像保存為 astronaut_rides_horse.png 文件。

若要在本地端跑，需確保是有網路的，才可以從Hugging Face 下載並載入模型喔