Update app.py

app.py

@@ -1,50 +1,86 @@
 import torch
 from diffusers import AutoPipelineForText2Image
-from peft import PeftModel, PeftConfig
-import gradio as gr # Import Gradio
+import gradio as gr
 
-# 기기 설정
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# เลือก device เป็น cuda ถ้ามี GPU เพื่อให้ offload ทำงานร่วมกับ GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 기본 모델 로드
-print("기본 FLUX 모델 로드 중...")
+print("โหลด FLUX model ด้วยการตั้งค่าสำหรับ offload...")
 pipe = AutoPipelineForText2Image.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.float16  # bfloat16 대신 float16 사용
+    torch_dtype=torch.float16  # ลด precision เพื่อประหยัดหน่วยความจำ
 )
-pipe.to(device)
 
-# Uncensored LoRA 로드
-print("Uncensored LoRA 로드 중...")
+# ถ้ามี GPU ให้ย้าย pipeline ไปยัง GPU ก่อนเปิด offload (diffusers จะจัดการย้ายพารามิเตอร์)
+if device == "cuda":
+    pipe.to("cuda")
+
+# เปิดการประหยัดหน่วยความจำ
+# 1) attention slicing ลด peak memory ขณะ attention
+pipe.enable_attention_slicing()
+# 2) VAE slicing ลด memory ตอน decode
+try:
+    pipe.enable_vae_slicing()
+except Exception:
+    pass
+
+# 3) เปิด CPU offload: ย้ายพารามิเตอร์โมเดลไปเก็บใน CPU RAM แล้วโหลดไป GPU เฉพาะตอนใช้งาน
+#    ถ้าต้องการให้ใช้โฟลเดอร์สำหรับ offload ชั่วคราว ให้กำหนด offload_folder
+try:
+    pipe.enable_model_cpu_offload(
+        gpu_id=0,                # หมายเลข GPU ที่จะใช้ (ถ้ามีหลายตัว ปรับตามจริง)
+        offload_folder="./offload",  # ถ้าต้องการใช้ disk เป็นสำรอง
+        pin_memory=True          # pin memory ช่วยให้ย้ายข้อมูลเร็วขึ้น
+    )
+except Exception:
+    # ถ้าเวอร์ชัน diffusers ไม่มีฟังก์ชันนี้ ให้ลอง enable_sequential_cpu_offload
+    try:
+        pipe.enable_sequential_cpu_offload()
+    except Exception:
+        pass
+
+# เพิ่มการตั้งค่าอื่น ๆ ที่ช่วยลด memory
+torch.backends.cudnn.benchmark = False
+
+# โหลด LoRA (ถ้าจำเป็น) — ควรโหลดหลังจาก pipeline ถูกตั้งค่า offload แล้ว
+print("โหลด Uncensored LoRA...")
 pipe.load_lora_weights(
     'Heartsync/Flux-NSFW-uncensored',
     weight_name='lora.safetensors',
     adapter_name="uncensored"
 )
 
-# 이미지 생성 함수 정의
+# ฟังก์ชันสร้างภาพ
 def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, width, height, seed):
-    generator = torch.Generator(device=device).manual_seed(seed)
+    # แนะนำ: ถ้าต้องการจำกัด VRAM ให้ลด width/height หรือใช้ smaller batch
+    generator = torch.Generator(device="cpu").manual_seed(int(seed))
 
-    image = pipe(
+    # ถ้า GPU มี ให้ใช้ device="cuda" ในการเรียก แต่ pipeline จะจัดการ offload ให้
+    out = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
+        guidance_scale=float(guidance_scale),
+        num_inference_steps=int(num_inference_steps),
+        width=int(width),
+        height=int(height),
         generator=generator,
-    ).images[0]
-    
-    # 이미지 저장 (선택 사항, 필요에 따라 활성화/비활성화)
-    # image.save("generated_image.png")
+    )
+    image = out.images[0]
+
+    # เคลียร์ cache ของ CUDA ถ้ามี
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+
     return image
 
-# Gradio 인터페이스 생성
+# Gradio interface (เหมือนเดิม)
 iface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.Textbox(label="Prompt", value="A woman in a sheer white dress standing on a beach at sunset, backlit so her silhouette is visible through the thin fabric, shot with Canon EOS R5, 85mm f/1.2 lens, golden hour natural lighting, professional composition, hyperrealistic detail, masterpiece quality, 8K resolution."),
+        gr.Textbox(label="Prompt", value="A woman in a sheer white dress standing on a beach at sunset..."),
         gr.Textbox(label="Negative Prompt", value="text, watermark, signature, cartoon, anime, illustration, painting, drawing, low quality, blurry"),
         gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.0, label="Guidance Scale"),
         gr.Slider(minimum=10, maximum=100, step=1, value=28, label="Number of Inference Steps"),
@@ -53,9 +89,8 @@ iface = gr.Interface(
         gr.Slider(minimum=0, maximum=99999, step=1, value=42, label="Seed")
     ],
     outputs="image",
-    title="FLUX.1-dev with Uncensored LoRA",
-    description="Generate images using FLUX.1-dev with a loaded Uncensored LoRA model."
+    title="FLUX.1-dev with Uncensored LoRA (Offload to RAM)",
+    description="Generate images using FLUX.1-dev with CPU offload to keep GPU VRAM low."
 )
 
-# 인터페이스 실행 (share=True)
 iface.launch(share=True)