Update app.py

app.py

@@ -12,6 +12,9 @@ import torch
 from diffusers import DiffusionPipeline
 from transformers import pipeline
 import gradio as gr
+import os
+from huggingface_hub import login
+
 
 # Download checkpoints
 snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
@@ -38,6 +41,27 @@ pt_model = LeffaModel(
 )
 pt_inference = LeffaInference(model=pt_model)
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("Please set the HF_TOKEN environment variable")
+login(token=HF_TOKEN)
+
+# FLUX 모델 초기화 부분 수정
+fashion_pipe = DiffusionPipeline.from_pretrained(
+    base_model,
+    torch_dtype=torch.bfloat16,
+    use_auth_token=HF_TOKEN  # 인증 토큰 추가
+)
+fashion_pipe.to("cuda")
+
+# LoRA 로딩 함수 수정
+def load_lora(pipe, repo_id):
+    pipe.load_lora_weights(
+        repo_id,
+        use_auth_token=HF_TOKEN
+    )
+    return pipe
+    
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 base_model = "black-forest-labs/FLUX.1-dev"
 model_lora_repo = "Motas/Flux_Fashion_Photography_Style"
@@ -58,12 +82,12 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
     else:
         actual_prompt = prompt
 
-    # 모드에 따른 LoRA 및 트리거워드 설정
+    # 모드에 따른 LoRA 로딩 및 트리거워드 설정
     if mode == "Generate Model":
-        pipe.load_lora_weights(model_lora_repo)
+        fashion_pipe = load_lora(fashion_pipe, model_lora_repo)
         trigger_word = "fashion photography, professional model"
     else:
-        pipe.load_lora_weights(clothes_lora_repo)
+        fashion_pipe = load_lora(fashion_pipe, clothes_lora_repo)
         trigger_word = "upper clothing, fashion item"
 
     if randomize_seed:
@@ -76,13 +100,14 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
         if i % (steps // 10) == 0:
             progress(i / steps * 100, f"Processing step {i} of {steps}...")
 
-    image = pipe(
+    image = fashion_pipe(
         prompt=f"{actual_prompt} {trigger_word}",
         num_inference_steps=steps,
         guidance_scale=cfg_scale,
         width=width,
         height=height,
         generator=generator,
+        use_auth_token=HF_TOKEN,  # 인증 토큰 추가
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]