Cho phép sinh nhieu anh

app.py

@@ -2,189 +2,151 @@ import gradio as gr
 import torch
 from transformers import pipeline, AutoModelForSeq2SeqLM, BartTokenizer, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline
-import io
 from PIL import Image
-import traceback
+import io
 import os
-from pathlib import Path
+import zipfile
+import traceback
 
-# === Thiết lập môi trường ===
+# === Thiết lập thiết bị ===
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Device: {device}")
 
-# === Load models với xử lý lỗi ===
-try:
-    # Summarizer (BART)
-    model_name = "lacos03/bart-base-finetuned-xsum"
-    print(f"Loading BART model from {model_name}...")
-    tokenizer = BartTokenizer.from_pretrained(model_name, use_fast=False)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
-    model.to(device)
-    summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=device)
-    print("✅ BART loaded successfully")
-except Exception as e:
-    print(f"❌ Error loading BART: {e}")
-    summarizer = None
-
-try:
-    # Promptist
-    print("Loading Promptist model...")
-    def load_prompter():
-        prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist", torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        tokenizer.pad_token = tokenizer.eos_token
-        tokenizer.padding_side = "left"
-        return prompter_model, tokenizer
-    promptist_model, promptist_tokenizer = load_prompter()
-    print("✅ Promptist loaded successfully")
-except Exception as e:
-    print(f"❌ Error loading Promptist: {e}")
-    promptist_model = None
-    promptist_tokenizer = None
-
-try:
-    # Stable Diffusion + LoRA
-    print("Loading Stable Diffusion model...")
-    sd_model_id = "runwayml/stable-diffusion-v1-5"
-    image_generator = StableDiffusionPipeline.from_pretrained(
-        sd_model_id,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        use_safetensors=True
-    ).to(device)
-    lora_weights = "lacos03/std-1.5-lora-midjourney-1.0"
-    print(f"Loading LoRA weights from {lora_weights}...")
-    image_generator.load_lora_weights(lora_weights)
-    print("✅ Stable Diffusion with LoRA loaded successfully")
-except Exception as e:
-    print(f"❌ Error loading Stable Diffusion or LoRA: {e}")
-    image_generator = None
-
-# === Modular hóa ===
-def summarize(article_text):
-    if not summarizer or not article_text.strip():
-        return "[Empty input or model not loaded]", "[Empty input or model not loaded]"
-    try:
-        summary = summarizer(article_text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-        title = summary.split(".")[0] + "."
-        return title, summary
-    except Exception as e:
-        return f"[Error in summarization: {e}]", f"[Error in summarization: {e}]"
-
-def generate_prompt(title):
-    if not promptist_model or not promptist_tokenizer or not title:
-        return "[Error: Promptist not loaded or no title]"
-    try:
-        input_ids = promptist_tokenizer(title.strip() + " Rephrase:", return_tensors="pt").input_ids.to(device)
-        eos_id = promptist_tokenizer.eos_token_id
-        outputs = promptist_model.generate(
-            input_ids,
-            do_sample=False,
-            max_new_tokens=75,
-            num_beams=8,
-            num_return_sequences=8,
-            eos_token_id=eos_id,
-            pad_token_id=eos_id,
-            length_penalty=-1.0
-        )
-        output_texts = promptist_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        prompt = output_texts[0].replace(title + " Rephrase:", "").strip()
-        return prompt
-    except Exception as e:
-        return f"[Error in prompt generation: {e}]"
-
-def generate_image(prompt, style):
-    if not image_generator or not prompt:
-        blank = Image.new("RGB", (512, 512), (255, 255, 255))
-        img_byte_arr = io.BytesIO()
-        blank.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
-        return blank, img_byte_arr
-    try:
-        styled_prompt = f"{prompt}, {style.lower()} style"
-        result = image_generator(
-            styled_prompt,
-            num_inference_steps=50,
-            guidance_scale=7.5
-        ).images[0]
-        img_byte_arr = io.BytesIO()
-        result.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
-        return result, img_byte_arr
-    except Exception as e:
-        print(f"❌ Image generation error: {traceback.format_exc()}")
-        blank = Image.new("RGB", (512, 512), (255, 255, 255))
-        img_byte_arr = io.BytesIO()
-        blank.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
-        return blank, img_byte_arr
-
-# === Main processing function with staged outputs ===
-def process_step_by_step(article_text, style_choice, state=None):
-    if state is None:
-        state = {"title": None, "prompt": None, "image": None, "file_path": None}
-    
-    # Bước 1: Tóm tắt và tạo tiêu đề
-    title, summary = summarize(article_text)
-    print(f"Summary title: {title}")
-    state["title"] = title
-    yield state, title, None, None, None
-    
-    # Bước 2: Tạo prompt
-    prompt = generate_prompt(title)
-    print(f"Generated prompt: {prompt}")
-    state["prompt"] = prompt
-    yield state, title, prompt, None, None
-    
-    # Bước 3: Tạo ảnh
-    image, img_bytes = generate_image(prompt, style_choice)
-    print(f"Image generated: {image.size if image else 'None'}")
-    
-    # Lưu ảnh tạm thời
-    temp_dir = "./temp"
+# === Load models ===
+# BART Summarizer
+model_name = "lacos03/bart-base-finetuned-xsum"
+tokenizer = BartTokenizer.from_pretrained(model_name, use_fast=False)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
+summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=0 if device=="cuda" else -1)
+
+# Promptist
+promptist_model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Promptist",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
+promptist_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+promptist_tokenizer.pad_token = promptist_tokenizer.eos_token
+promptist_tokenizer.padding_side = "left"
+
+# Stable Diffusion + LoRA
+sd_model_id = "runwayml/stable-diffusion-v1-5"
+image_generator = StableDiffusionPipeline.from_pretrained(
+    sd_model_id,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    use_safetensors=True
+).to(device)
+lora_weights = "lacos03/std-1.5-lora-midjourney-1.0"
+image_generator.load_lora_weights(lora_weights)
+
+# === Hàm xử lý ===
+def summarize_article(article_text):
+    """Tóm tắt bài viết và tạo prompt refinement"""
+    if not article_text.strip():
+        return "[Empty]", "[Empty]"
+    summary = summarizer(article_text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
+    title = summary.split(".")[0] + "."
+    # Prompt refinement
+    input_ids = promptist_tokenizer(title.strip() + " Rephrase:", return_tensors="pt").input_ids.to(device)
+    eos_id = promptist_tokenizer.eos_token_id
+    outputs = promptist_model.generate(
+        input_ids,
+        do_sample=False,
+        max_new_tokens=75,
+        num_beams=8,
+        num_return_sequences=1,
+        eos_token_id=eos_id,
+        pad_token_id=eos_id,
+        length_penalty=-1.0
+    )
+    output_texts = promptist_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    prompt = output_texts[0].replace(title + " Rephrase:", "").strip()
+    return title, prompt
+
+def generate_images(prompt, style, num_images=4):
+    """Sinh nhiều ảnh"""
+    styled_prompt = f"{prompt}, {style.lower()} style"
+    results = image_generator(
+        styled_prompt,
+        num_inference_steps=50,
+        guidance_scale=7.5,
+        num_images_per_prompt=num_images
+    ).images
+    return results
+
+def save_selected_images(selected_idx, all_images):
+    """Lưu ảnh đã chọn và nén thành ZIP"""
+    if not selected_idx:
+        return None
+    temp_dir = "./temp_selected"
     os.makedirs(temp_dir, exist_ok=True)
-    temp_file = os.path.join(temp_dir, f"generated_image_{id(image)}.png")
-    image.save(temp_file, format="PNG")
-    state["image"] = image
-    state["file_path"] = temp_file
-    
-    print(f"✅ Process completed")
-    yield state, title, prompt, image, temp_file
-
-# === Gradio UI ===
+    zip_path = os.path.join(temp_dir, "selected_images.zip")
+    with zipfile.ZipFile(zip_path, 'w') as zipf:
+        for idx in selected_idx:
+            img = all_images[int(idx)]
+            img_path = os.path.join(temp_dir, f"image_{idx}.png")
+            img.save(img_path, format="PNG")
+            zipf.write(img_path, f"image_{idx}.png")
+    return zip_path
+
+# === UI Gradio ===
 def create_app():
     with gr.Blocks() as demo:
-        gr.Markdown("## 📰 Article → 🖼️ Image Generator")
-        gr.Markdown("Nhập bài viết → Sinh tiêu đề → Tối ưu prompt → Sinh ảnh minh họa tự động")
-
-        # State để lưu trữ trạng thái tạm thời
-        state = gr.State(value=None)
+        gr.Markdown("## 📰 Article → 🖼️ Multiple Image Generator with Selection")
 
+        # Bước 1: Nhập bài viết và sinh tiêu đề + prompt
         with gr.Row():
-            article_input = gr.Textbox(label="📄 Bài viết", lines=10, placeholder="Dán nội dung bài viết ở đây...")
-            style_dropdown = gr.Dropdown(choices=["Art", "Anime", "Watercolor", "Cyberpunk"], label="🎨 Phong cách ảnh", value="Art")
+            article_input = gr.Textbox(label="📄 Bài viết", lines=10)
+            style_dropdown = gr.Dropdown(
+                choices=["Art", "Anime", "Watercolor", "Cyberpunk"],
+                label="🎨 Phong cách ảnh", value="Art"
+            )
+        num_images_slider = gr.Slider(1, 8, value=4, step=1, label="🔢 Số lượng ảnh")
+        btn_summary = gr.Button("📌 Sinh Tiêu đề & Prompt")
+
+        title_output = gr.Textbox(label="Tiêu đề")
+        prompt_output = gr.Textbox(label="Prompt sinh ảnh")
+
+        # Bước 2: Sinh ảnh từ prompt đã refine
+        btn_generate_images = gr.Button("🎨 Sinh ảnh từ Prompt")
+        gallery = gr.Gallery(label="🖼️ Ảnh minh họa", columns=2, height=600)
+        selected_indices = gr.CheckboxGroup(choices=[], label="Chọn ảnh để tải về")
+
+        # Bước 3: Tải ảnh đã chọn
+        btn_download = gr.Button("📥 Tải ảnh đã chọn")
+        download_file = gr.File(label="File ZIP tải về")
+
+        # Logic
+        btn_summary.click(
+            fn=summarize_article,
+            inputs=[article_input],
+            outputs=[title_output, prompt_output]
+        )
 
-        with gr.Row():
-            submit_button = gr.Button("🚀 Tạo Tiêu đề & Ảnh Minh họa")
+        def update_gallery(prompt, style, num_images):
+            images = generate_images(prompt, style, num_images)
+            choices = [str(i) for i in range(len(images))]
+            return images, gr.update(choices=choices, value=[]), images  # images lưu tạm trong state
 
-        with gr.Row():
-            title_output = gr.Textbox(label="📌 Tiêu đề được tạo")
-            prompt_output = gr.Textbox(label="🔧 Prompt sinh ảnh")
+        image_state = gr.State([])
 
-        image_output = gr.Image(label="🖼️ Ảnh minh họa", interactive=True)
-        download_button = gr.File(label="📥 Tải ảnh")
-
-        feedback = gr.Radio(["👍 Hài lòng", "👎 Không hài lòng"], label="📊 Bạn có hài lòng với kết quả không?", value=None)
+        btn_generate_images.click(
+            fn=update_gallery,
+            inputs=[prompt_output, style_dropdown, num_images_slider],
+            outputs=[gallery, selected_indices, image_state]
+        )
 
-        # Gắn sự kiện nút submit với hàm process từng bước
-        submit_button.click(
-            fn=process_step_by_step,
-            inputs=[article_input, style_dropdown, state],
-            outputs=[state, title_output, prompt_output, image_output, download_button]
+        btn_download.click(
+            fn=save_selected_images,
+            inputs=[selected_indices, image_state],
+            outputs=[download_file]
         )
 
     return demo
 
-# === Launch ===
+# === Chạy app ===
 if __name__ == "__main__":
     app = create_app()
-    app.launch(debug=True, share=True)
+    app.launch(debug=True, share=True)