Update app.py

app.py

@@ -722,6 +722,34 @@ def save_generated_image(image, prompt):
     image.save(filepath)
     return filepath
 
+gen_pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.float16,
+    use_auth_token=HF_TOKEN
+)
+gen_pipe.enable_attention_slicing(slice_size="auto")
+
+# 이미지 생성용 LoRA 가중치 로드
+try:
+    lora_path = hf_hub_download(
+        "ginipick/flux-lora-eric-cat",
+        "flux-lora-eric-cat.safetensors",  # 실제 파일명 확인 필요
+        use_auth_token=HF_TOKEN
+    )
+    gen_pipe.load_lora_weights(lora_path)
+    gen_pipe.fuse_lora(lora_scale=0.125)
+except Exception as e:
+    print(f"Error loading generation LoRA weights: {str(e)}")
+    raise ValueError("Failed to load generation LoRA weights. Please check your HF_TOKEN and model access.")
+
+# GPU로 이동
+if torch.cuda.is_available():
+    try:
+        gen_pipe = gen_pipe.to("cuda:0")
+    except Exception as e:
+        print(f"Warning: Could not move generation pipeline to CUDA: {str(e)}")
+
+# generate_image 함수 수정
 @spaces.GPU(duration=60)
 def generate_image(
     prompt: str,
@@ -740,7 +768,8 @@ def generate_image(
         generator = torch.Generator(device=device).manual_seed(seed)
         
         with torch.inference_mode():
-            image = pipe(
+            # gen_pipe 사용
+            image = gen_pipe(
                 prompt=prompt,
                 width=width,
                 height=height,
@@ -757,16 +786,13 @@ def generate_image(
     finally:
         clear_memory()
 
-
-
-
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     position = gr.State(value="bottom-center")
 
     gr.HTML("""
         <div class="main-title">
-            <h1>🎨 GiniGen Canvas-o3</h1>
-            <p>Remove background of specified objects, generate new backgrounds, and insert text over or behind images with prompts.</p>
+            <h1>🎨 Webtoon Canvas</h1>
+            <p>Webtoon generated, Remove background of specified objects, generate new backgrounds, and insert text over or behind images with prompts.</p>
         </div>
     """)