Update app.py

app.py

@@ -27,7 +27,40 @@ from selenium.common.exceptions import WebDriverException, TimeoutException
 from PIL import Image
 from io import BytesIO
 from datetime import datetime
-
+import spaces
+from safetensors.torch import load_file
+from diffusers import FluxPipeline
+import torch
+
+# 캐시 경로 설정
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+# FLUX 모델 초기화
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+# 이미지 생성 함수 추가
+@spaces.GPU
+def generate_image(prompt, height=512, width=512, steps=8, scales=3.5, seed=3413):
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        return pipe(
+            prompt=[prompt],
+            generator=torch.Generator().manual_seed(int(seed)),
+            num_inference_steps=int(steps),
+            guidance_scale=float(scales),
+            height=int(height),
+            width=int(width),
+            max_sequence_length=256
+        ).images[0]
+        
 # SystemPrompt 부분을 직접 정의
 SystemPrompt = """You are 'MOUSE-I', an advanced AI visualization expert. Your mission is to transform every response into a visually stunning and highly informative presentation.