Spaces:
Sleeping
Sleeping
| 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 | |
| # --- Management Section Cache: Bring the NSFW model to the standard model. --- | |
| MODEL_STANDARD = "MiaoshouAI/Florence-2-large-PromptGen-v2.0" | |
| MODEL_NSFW = "ljnlonoljpiljm/florence-2-large-nsfw-pt" | |
| def setup_model_cache(): | |
| """ Download Model NSFW Let's place the folder of the standard model in Cache """ | |
| cache_dir = os.path.expanduser("~/.cache/huggingface/hub") | |
| 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"🔍 Check Cache: {folder_standard}") | |
| if not os.path.exists(path_standard) or not os.listdir(path_standard): | |
| print(f"⚠️ No standard models found in Cache or empty folders") | |
| if os.path.exists(path_nsfw) and os.listdir(path_nsfw): | |
| print(f"✅ NSFW model has been found in Cache.: {folder_nsfw}") | |
| source_path = path_nsfw | |
| else: | |
| print(f"🚀 Downloading Model NSFW ({MODEL_NSFW})...") | |
| try: | |
| snapshot_download( | |
| repo_id=MODEL_NSFW, | |
| local_dir=path_nsfw, | |
| local_dir_use_symlinks=False | |
| ) | |
| print("✅ Download Model NSFW Finished") | |
| source_path = path_nsfw | |
| except Exception as e: | |
| print(f"❌ Download failed: {e}") | |
| print("💡 Use standard models instead (may not have NSFW filter)") | |
| return MODEL_STANDARD | |
| if os.path.exists(path_standard): | |
| shutil.rmtree(path_standard) | |
| print(f"📂 Power Copy File from {folder_nsfw} -> {folder_standard}...") | |
| shutil.copytree(source_path, path_standard) | |
| print("✅ Overlay file finish!") | |
| else: | |
| print(f"✅ Models have been found in Cache.: {folder_standard}") | |
| return MODEL_STANDARD | |
| # --- เรียกฟังก์ชันจัดการ Cache --- | |
| FINAL_MODEL_NAME = setup_model_cache() | |
| print(f"🚀 Loading model (customized): {FINAL_MODEL_NAME}...") | |
| # --- ติดตั้ง flash-attn (ถ้าจำเป็น) --- | |
| try: | |
| import flash_attn | |
| print("✅ flash_attn Available") | |
| except ImportError: | |
| print("⚠️ flash_attn not Available...") | |
| 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("✅ Load Model Florence-2 (NSFW Version) Finished!") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| florence_model = None | |
| florence_processor = None | |
| # --- รายการ Task ที่รองรับ (As you wish?) --- | |
| TASK_PROMPTS = { | |
| "Analyze (only PromptGen models)": "<ANALYZE>", | |
| "Generate tags (PromptGen only)": "<GENERATE_TAGS>", | |
| "Caption & Tags (PromptGen only)": "<MIXED_CAPTION>", | |
| "Caption & Tags++ (PromptGen only)": "<MIXED_CAPTION_PLUS>", | |
| "Caption": "<CAPTION>", | |
| "Detailed Caption": "<DETAILED_CAPTION>", | |
| "More Detailed Caption": "<MORE_DETAILED_CAPTION>", | |
| "OCR": "<OCR>", | |
| "OCR with Region": "<OCR_WITH_REGION>", | |
| "Object Detection": "<OD>", | |
| "Dense Region Caption": "<DENSE_REGION_CAPTION>", | |
| "Region Proposal": "<REGION_PROPOSAL>", | |
| "Caption to Phrase Grounding": "<CAPTION_TO_PHRASE_GROUNDING>", | |
| "Referring Expression Segmentation": "<REFERRING_EXPRESSION_SEGMENTATION>", | |
| "Region to Segmentation": "<REGION_TO_SEGMENTATION>", | |
| "Open Vocabulary Detection": "<OPEN_VOCABULARY_DETECTION>", | |
| "Region to Category": "<REGION_TO_CATEGORY>", | |
| "Region to Description": "<REGION_TO_DESCRIPTION>", | |
| } | |
| def process_image(image, task_name): | |
| global florence_model, florence_processor | |
| if florence_model is None or florence_processor is None: | |
| return "❌ The model has not been loaded or there is an error in the start." | |
| else: | |
| print(f"Model: {florence_model}\nProcessor: {florence_processor}") | |
| if image is None: | |
| return "Please select a picture" | |
| try: | |
| # แปลง Gradio Image เป็น PIL Image | |
| if not isinstance(image, Image.Image): | |
| image = Image.fromarray(image) | |
| if image.mode != "RGB": | |
| image = image.convert("RGB") | |
| # ดึง Prompt จาก Task ที่เลือก | |
| prompt = TASK_PROMPTS.get(task_name, "<CAPTION>") | |
| inputs = florence_processor(text=prompt, 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=prompt, image_size=(image.width, image.height) | |
| ) | |
| # --- ส่วนแก้ไข: แยกค่าออกจาก Dictionary --- | |
| # parsed_answer จะเป็น dict เช่น {'<DETAILED_CAPTION>': 'ข้อความ...'} | |
| # เราต้องการแค่ 'ข้อความ...' | |
| if isinstance(parsed_answer, dict): | |
| # ดึงค่าแรก (Value) จาก Dictionary | |
| result_text = next(iter(parsed_answer.values())) | |
| else: | |
| # ถ้าไม่ใช่ dict (กรณีบาง task อาจ return string โดยตรง) | |
| result_text = str(parsed_answer) | |
| print(f"\n✅ Task: {task_name} | Result: {result_text}") | |
| return result_text | |
| except Exception as e: | |
| return f"❌ Error while processing: {str(e)}" | |
| # --- สร้าง UI --- | |
| with gr.Blocks(title="Image-to-Prompt (Florence-2)", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🖼️ Image-to-Prompt (Florence-2) | |
| Upload the image and select the analysis type to create Prompt | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image(type="pil", label="Pictures", height=400) | |
| task_dropdown = gr.Dropdown( | |
| choices=list(TASK_PROMPTS.keys()), | |
| value="More Detailed Caption", # เลือกเป็นค่าเริ่มต้น | |
| label="Analysis Type" | |
| ) | |
| btn = gr.Button("🚀 Create results", variant="primary") | |
| with gr.Column(): | |
| text_output = gr.Textbox(label="ผลลัพธ์", lines=10, max_lines=20, show_copy_button=True) | |
| btn.click( | |
| fn=process_image, | |
| inputs=[image_input, task_dropdown], | |
| outputs=text_output | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) |