Fix

app.py

@@ -1,57 +1,105 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForSeq2SeqLM, BartTokenizer, GenerationConfig, AutoModelForCausalLM, AutoTokenizer
-from diffusers import StableDiffusionPipeline
 import torch
+from transformers import pipeline, AutoModelForSeq2SeqLM, BartTokenizer, AutoModelForCausalLM, AutoTokenizer
+from diffusers import StableDiffusionPipeline
 import io
 from PIL import Image
 import traceback
+import os
+from pathlib import Path
 
-# === Load models ===
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Summarizer (BART)
-model_name = "lacos03/bart-base-finetuned-xsum"
-tokenizer = BartTokenizer.from_pretrained(model_name, use_fast=False)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-generation_config = GenerationConfig.from_pretrained(model_name)
-generation_config.early_stopping = True
-summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, generation_config=generation_config)
-
-# Promptist
-promptist_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
-promptist_tokenizer = AutoTokenizer.from_pretrained("microsoft/Promptist")
-
-# 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)
+# === Thiết lập môi trường ===
+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:
-        if not article_text.strip():
-            return "[Empty input]", "[Empty input]"
         summary = summarizer(article_text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-        title = summary.split(".")[0]
+        title = summary.split(".")[0] + "."
         return title, summary
     except Exception as e:
-        return "[Error in summarization]", str(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:
-        inputs = promptist_tokenizer(title, return_tensors="pt").to(device)
-        output = promptist_model.generate(**inputs, max_length=50, num_return_sequences=1)
-        prompt = promptist_tokenizer.decode(output[0], skip_special_tokens=True)
+        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 "[Error in prompt generation]"
+        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(
@@ -64,16 +112,35 @@ def generate_image(prompt, style):
         img_byte_arr.seek(0)
         return result, img_byte_arr
     except Exception as e:
-        print(traceback.format_exc())
+        print(f"❌ Image generation error: {traceback.format_exc()}")
         blank = Image.new("RGB", (512, 512), (255, 255, 255))
-        return blank, io.BytesIO()
+        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 ===
 def process(article_text, style_choice):
+    print(f"Processing article: {article_text[:50]}...")
     title, summary = summarize(article_text)
+    print(f"Summary title: {title}")
     prompt = generate_prompt(title)
+    print(f"Generated prompt: {prompt}")
     image, img_bytes = generate_image(prompt, style_choice)
-    return title, prompt, image, img_bytes
+    print(f"Image generated: {image.size if image else 'None'}")
+
+    # Chuyển BytesIO thành file tạm và trả về đường dẫn
+    temp_dir = "./temp"
+    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")
+    with open(temp_file, "rb") as f:
+        img_file = f.read()
+    # Trả về đường dẫn tạm thời cho Gradio
+    file_path = temp_file
+
+    print(f"✅ Process completed")
+    return title, prompt, image, file_path
 
 # === Gradio UI ===
 def create_app():
@@ -83,7 +150,7 @@ def create_app():
 
         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=["Realistic", "Anime", "Watercolor", "Cyberpunk"], label="🎨 Phong cách ảnh", value="Realistic")
+            style_dropdown = gr.Dropdown(choices=["Art", "Anime", "Watercolor", "Cyberpunk"], label="🎨 Phong cách ảnh", value="Art")
 
         with gr.Row():
             submit_button = gr.Button("🚀 Tạo Tiêu đề & Ảnh Minh họa")
@@ -91,7 +158,7 @@ def create_app():
         with gr.Row():
             title_output = gr.Textbox(label="📌 Tiêu đề được tạo")
             prompt_output = gr.Textbox(label="🔧 Prompt sinh ảnh")
-        
+
         image_output = gr.Image(label="🖼️ Ảnh minh họa", interactive=True)
         download_button = gr.File(label="📥 Tải ảnh")
 
@@ -108,4 +175,4 @@ def create_app():
 # === Launch ===
 if __name__ == "__main__":
     app = create_app()
-    app.launch()
+    app.launch(debug=True, share=True)