Rename app.py.bak to app.py

app.py

@@ -0,0 +1,148 @@
+import gradio as gr
+import subprocess
+import torch
+import os
+import shutil
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM
+from huggingface_hub import snapshot_download
+
+# --- ส่วนจัดการ Cache: ดึงโมเดล NSFW มาวางทับ Standard Model ---
+MODEL_STANDARD = "microsoft/Florence-2-base"
+MODEL_NSFW = "ljnlonoljpiljm/florence-2-base-nsfw-v2"
+
+def setup_model_cache():
+    """
+    ดาวน์โหลดโมเดล NSFW มาวางทับโฟลเดอร์ของโมเดลมาตรฐานใน Cache
+    เพื่อให้ from_pretrained(MODEL_STANDARD) โหลดไฟล์ของ NSFW แทน
+    """
+    cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
+    
+    # ชื่อโฟลเดอร์ใน cache (แปลง / เป็น --)
+    folder_standard = f"models--{MODEL_STANDARD.replace('/', '--')}"
+    folder_nsfw = f"models--{MODEL_NSFW.replace('/', '--')}"
+    
+    path_standard = os.path.join(cache_dir, folder_standard)
+    path_nsfw = os.path.join(cache_dir, folder_nsfw)
+    
+    print(f"🔍 ตรวจสอบ Cache: {folder_standard}")
+    
+    # ถ้าโฟลเดอร์มาตรฐานยังไม่มี หรือต้องการบังคับอัปเดต
+    if not os.path.exists(path_standard) or not os.listdir(path_standard):
+        print(f"⚠️ ไม่พบโมเดลมาตรฐานใน Cache หรือโฟลเดอร์ว่าง")
+        
+        # ถ้ามี NSFW แล้ว ก็ไม่ต้องโหลดใหม่
+        if os.path.exists(path_nsfw) and os.listdir(path_nsfw):
+            print(f"✅ พบโมเดล NSFW ใน Cache แล้ว: {folder_nsfw}")
+            source_path = path_nsfw
+        else:
+            print(f"🚀 กำลังดาวน์โหลดโมเดล NSFW ({MODEL_NSFW})...")
+            try:
+                # ดาวน์โหลดลงโฟลเดอร์ชั่วคราว (ใช้ชื่อ NSFW)
+                snapshot_download(
+                    repo_id=MODEL_NSFW,
+                    local_dir=path_nsfw,
+                    local_dir_use_symlinks=False
+                )
+                print("✅ ดาวน์โหลดโมเดล NSFW เสร็จสิ้น")
+                source_path = path_nsfw
+            except Exception as e:
+                print(f"❌ ดาวน์โหลดล้มเหลว: {e}")
+                print("💡 ใช้โมเดลมาตรฐานแทน (อาจไม่มี NSFW filter)")
+                return MODEL_STANDARD # คืนชื่อมาตรฐานให้โหลดปกติ
+        
+        # ลบโฟลเดอร์มาตรฐานเดิม (ถ้ามี)
+        if os.path.exists(path_standard):
+            print(f"🗑️ ลบโฟลเดอร์มาตรฐานเดิม: {folder_standard}")
+            shutil.rmtree(path_standard)
+            
+        # Copy ไฟล์ทั้งหมดจาก NSFW มาทับ Standard
+        print(f"📂 กำลัง Copy ไฟล์จาก {folder_nsfw} -> {folder_standard}...")
+        shutil.copytree(source_path, path_standard)
+        print("✅ วางไฟล์ทับเสร็จสิ้น! ตอนนี้ from_pretrained('microsoft/Florence-2-base') จะใช้ไฟล์ของ NSFW")
+        
+    else:
+        print(f"✅ พบโมเดลใน Cache แล้ว: {folder_standard}")
+        # ตรวจสอบว่าไฟล์ภายในเป็นของ NSFW หรือไม่ (โดยดูไฟล์ config.json ถ้ามี)
+        # หากต้องการบังคับอัปเดตทุกครั้ง สามารถลบโฟลเดอร์นี้ทิ้งได้
+        
+    return MODEL_STANDARD
+
+# --- เรียกฟังก์ชันจัดการ Cache ก่อนโหลดโมเดล ---
+FINAL_MODEL_NAME = setup_model_cache()
+print(f"🚀 กำลังโหลดโมเดล (ที่ถูกปรับแต่งแล้ว): {FINAL_MODEL_NAME}...")
+
+# ติดตั้ง flash-attn (ถ้าจำเป็น)
+try:
+    import flash_attn
+    print("✅ flash_attn พร้อมใช้งาน")
+except ImportError:
+    print("⚠️ flash_attn ไม่พบ กำลังติดตั้ง...")
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+# --- โหลดโมเดล ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+florence_model = None
+florence_processor = None
+
+try:
+    # โหลดโมเดล (ตอนนี้จะโหลดไฟล์ที่เราก๊อ��ปี้มาแทนที่)
+    florence_model = AutoModelForCausalLM.from_pretrained(
+        FINAL_MODEL_NAME, 
+        trust_remote_code=True
+    ).to(device).eval()
+    
+    florence_processor = AutoProcessor.from_pretrained(
+        FINAL_MODEL_NAME, 
+        trust_remote_code=True
+    )
+    print("✅ โหลดโมเดล Florence-2 (NSFW Version) เสร็จสิ้น!")
+except Exception as e:
+    print(f"❌ เกิดข้อผิดพลาดในการโหลดโมเดล: {e}")
+    print("💡 ตรวจสอบ Logs เพื่อดูรายละเอียด")
+    florence_model = None
+    florence_processor = None
+
+def generate_caption(image):
+    global florence_model, florence_processor
+    
+    if florence_model is None or florence_processor is None:
+        return "❌ โมเดลยังไม่ได้โหลดหรือเกิดข้อผิดพลาดในการเริ่มต้น"
+        
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+        
+    try:
+        inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
+        generated_ids = florence_model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            early_stopping=False,
+            do_sample=False,
+            num_beams=3,
+        )
+        generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        parsed_answer = florence_processor.post_process_generation(
+            generated_text, 
+            task="<MORE_DETAILED_CAPTION>", 
+            image_size=(image.width, image.height)
+        )
+        prompt = parsed_answer["<MORE_DETAILED_CAPTION>"]
+        print("\n\nGeneration completed!:" + prompt)
+        return prompt
+    except Exception as e:
+        return f"❌ เกิดข้อผิดพลาดขณะประมวลผล: {str(e)}"
+
+# --- สร้าง UI ---
+io = gr.Interface(
+    fn=generate_caption, 
+    inputs=[gr.Image(label="Input Image")], 
+    outputs=[gr.Textbox(label="Output Prompt", lines=2, show_copy_button=True)], 
+    deep_link=False,
+    title="Image-to-Prompt (Florence-2 NSFW)",
+    description="อัปโหลดรูปภาพเพื่อสร้าง Prompt (ใช้โมเดล NSFW ที่ปรับแต่งแล้ว)"
+)
+
+if __name__ == "__main__":
+    io.launch(debug=True)

app.py.bak

@@ -1,69 +0,0 @@
-import gradio as gr
-import subprocess
-import torch
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
-
-# import os
-# import random
-# from gradio_client import Client
-
-
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
-# Initialize Florence model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
-
-# api_key = os.getenv("HF_READ_TOKEN")
-
-def generate_caption(image):
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    
-    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
-    generated_ids = florence_model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        early_stopping=False,
-        do_sample=False,
-        num_beams=3,
-    )
-    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = florence_processor.post_process_generation(
-        generated_text,
-        task="<MORE_DETAILED_CAPTION>",
-        image_size=(image.width, image.height)
-    )
-    prompt =  parsed_answer["<MORE_DETAILED_CAPTION>"]
-    print("\n\nGeneration completed!:"+ prompt)
-    return prompt
-    # yield prompt, None
-    # image_path = generate_image(prompt,random.randint(0, 4294967296))
-    # yield prompt, image_path 
-
-# def generate_image(prompt, seed=42, width=1024, height=1024):
-#     try:
-#         result = Client("KingNish/Realtime-FLUX", hf_token=api_key).predict(
-#             prompt=prompt,
-#             seed=seed,
-#             width=width,
-#             height=height,
-#             api_name="/generate_image"
-#         )
-#         # Extract the image path from the result tuple
-#         image_path = result[0]
-#         return image_path 
-#     except Exception as e:
-#         raise Exception(f"Error generating image: {str(e)}")
- 
-io = gr.Interface(generate_caption,
-                  inputs=[gr.Image(label="Input Image")],
-                  outputs = [gr.Textbox(label="Output Prompt", lines=2, show_copy_button = True),
-                             # gr.Image(label="Output Image")
-                            ],
-                  deep_link=False
-                 )
-io.launch(debug=True)