Update app.py

app.py

@@ -1,239 +1,369 @@
-import random
-import os
-import uuid
-from datetime import datetime
 import gradio as gr
 import numpy as np
-import spaces
-import torch
+import random
+import requests
+import io
+import re
+from datetime import datetime
+from PIL import Image, ImageDraw, ImageFont
+import textwrap
+import google.generativeai as genai
 from diffusers import DiffusionPipeline
-from PIL import Image
-from huggingface_hub import login
-
-# ---------- Authentication ----------
-hf_token = os.environ.get("HF_TOKEN")  # For HF Spaces or .env setups
-if hf_token:
-    login(token=hf_token)
-else:
-    print("⚠️ Warning: Hugging Face token not set. Cannot access gated models.")
-
-# ---------- 초기 설정 및 모델 로드 ----------
-SAVE_DIR = "saved_images"
-if not os.path.exists(SAVE_DIR):
-    os.makedirs(SAVE_DIR, exist_ok=True)
+import torch
 
+# Set device and model for SDXL-Turbo
 device = "cuda" if torch.cuda.is_available() else "cpu"
-repo_id = "black-forest-labs/FLUX.1-dev"
-adapter_id = "openfree/pepe"
+model_repo_id = "stabilityai/sdxl-turbo"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-pipeline.load_lora_weights(adapter_id)
-pipeline = pipeline.to(device)
+# Initialize SDXL-Turbo pipeline
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def save_generated_image(image, prompt):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    unique_id = str(uuid.uuid4())[:8]
-    filename = f"{timestamp}_{unique_id}.png"
-    filepath = os.path.join(SAVE_DIR, filename)
-    image.save(filepath)
-
-    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
-    with open(metadata_file, "a", encoding="utf-8") as f:
-        f.write(f"{filename}|{prompt}|{timestamp}\n")
-
-    return filepath
-
-def load_generated_images():
-    if not os.path.exists(SAVE_DIR):
-        return []
-    image_files = [
-        os.path.join(SAVE_DIR, f)
-        for f in os.listdir(SAVE_DIR)
-        if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))
-    ]
-    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
-    return image_files
-
-def load_predefined_images():
-    return []
+class CurrentAffairsGenerator:
+    def __init__(self, gemini_api_key, news_api_key=None, logo_url=None):
+        self.gemini_api_key = gemini_api_key
+        self.news_api_key = news_api_key
+        self.logo_url = logo_url
+        self.newsapi_base_url = "https://newsapi.org/v2/everything"
+        
+        if self.gemini_api_key:
+            try:
+                genai.configure(api_key=self.gemini_api_key)
+                self.gemini_model = genai.GenerativeModel('gemini-2.5-flash')
+            except Exception as e:
+                gr.Warning(f"Could not configure Gemini API: {e}.")
+                self.gemini_model = None
+        else:
+            self.gemini_model = None
+
+        self.brand_colors = {
+            'primary': '#0F172A',
+            'accent': '#F59E0B',
+            'text': '#FFFFFF',
+            'text_shadow': '#000000',
+        }
+
+    def fetch_latest_news(self, category):
+        if self.news_api_key:
+            gr.Info(f"Searching for live news about '{category}'...")
+            try:
+                news_data = self._search_with_newsapi(category)
+                if news_data:
+                    gr.Info("Successfully fetched live news!")
+                    return news_data
+                gr.Warning("Could not find relevant live news. Using a sample article.")
+            except Exception as e:
+                gr.Error(f"News API Error: {e}. Using a sample article.")
+        
+        return self._get_sample_news(category)
+
+    def _search_with_newsapi(self, category):
+        params = {
+            'q': f'"{category}" AND "India"',
+            'language': 'en',
+            'sortBy': 'publishedAt',
+            'pageSize': 10,
+            'apiKey': self.news_api_key
+        }
+        response = requests.get(self.newsapi_base_url, params=params)
+        response.raise_for_status()
+        articles = response.json().get('articles', [])
+        
+        for article in articles:
+            content = article.get('content')
+            description = article.get('description')
+            full_content = content if content and len(content) > 10 else description
+            
+            if full_content:
+                full_content_cleaned = re.sub(r'\[\s*\+\d+\s*chars\s*\]$', '', full_content).strip()
+                if len(full_content_cleaned) > 50:
+                    return {
+                        'title': article['title'],
+                        'full_content': full_content_cleaned,
+                        'source': article['source']['name'],
+                        'category': category
+                    }
+        return None
+
+    def _get_sample_news(self, topic=None):
+        sample_news = [
+            {'title': 'ISRO Confirms Gaganyaan Crew Module is Flight-Ready', 'full_content': 'The Indian Space Research Organisation has successfully completed all qualification tests for its first crewed space mission. This milestone includes rigorous vibration tests, acoustic assessments, and a series of parachute deployment simulations. The crew module is now officially certified for flight, bringing India one step closer to sending its own astronauts into space.', 'source': 'ISRO', 'category': 'Science'},
+            {'title': 'India Launches "Bharat AI" Research Supercluster', 'full_content': 'In a major push for artificial intelligence, the government has funded a new initiative to build a world-class AI ecosystem. This supercluster will connect top research institutions and startups, providing them with immense computing power and a platform for collaboration to solve India-specific challenges.', 'source': 'Ministry of Technology', 'category': 'Technology'},
+            {'title': 'UPI Transactions Cross 10 Billion Mark in a Single Month', 'full_content': 'India\'s Unified Payments Interface (UPI) has set a new global benchmark for digital payments, crossing the 10 billion transaction mark in a single month. This reflects the rapid and widespread adoption of digital finance across both urban and rural India, transforming the nation\'s economic landscape.', 'source': 'Reserve Bank of India', 'category': 'Economy'}
+        ]
+        if topic:
+            topic_lower = topic.lower()
+            filtered_news = [news for news in sample_news if topic_lower in news['category'].lower()]
+            if filtered_news: return random.choice(filtered_news)
+        return random.choice(sample_news)
+
+    def generate_llm_content(self, news_data):
+        if not self.gemini_model:
+            return news_data['title'], [f"Key Update: {news_data['full_content']}"]
+
+        prompt = f"""
+        Act as an expert social media analyst. Analyze this news:
+        Title: "{news_data['title']}"
+        Content: "{news_data['full_content']}"
+
+        Your Task:
+        1. Create one short, impactful headline (under 10 words).
+        2. Generate 3 highly informative bullet points.
+        3. IMPORTANT: For the bullet points, provide ONLY the text. DO NOT include any leading emojis or symbols like '☐', '-', or '*'.
+
+        Output Format (Strictly follow this):
+        HEADLINE: [Your headline here]
+        BULLETS:
+        [Your first insightful bullet point text]
+        [Your second insightful bullet point text]
+        [Your third insightful bullet point text]
+        """
+        try:
+            response = self.gemini_model.generate_content(prompt)
+            clean_text = re.sub(r'[\*#]', '', response.text)
+            
+            headline = clean_text.split("HEADLINE:")[1].split("BULLETS:")[0].strip()
+            bullets_str = clean_text.split("BULLETS:")[1].strip()
+            bullets = [line.strip().lstrip('- ') for line in bullets_str.split('\n') if line.strip()]
+            
+            return headline, bullets[:3]
+        except Exception as e:
+            gr.Warning(f"Gemini content generation failed: {e}. Using fallback.")
+            return news_data['title'], [news_data['full_content']]
+
+    def create_image_prompt(self, news_data):
+        category = news_data.get('category', 'General').lower()
+        prompts = {
+            'technology': [
+                "Futuristic cityscape with holographic data streams and glowing circuits, cinematic lighting, deep blue and gold palette, 4K hyper-detailed digital art.",
+                "Abstract visualization of a neural network, glowing nodes connected by light trails, dark background, symbolizing AI and connectivity.",
+                "Sleek server room with a liquid cooling system, neon blue lights reflecting on polished surfaces, representing high-performance computing.",
+                "Digital art of a human brain merging with a glowing circuit board, illustrating the concept of human-computer interface.",
+            ],
+            'economy': [
+                "Abstract visualization of a bullish stock market graph turning into a growing tree with coin-like leaves, against a modern city skyline.",
+                "A dynamic 3D render of the Indian Rupee symbol glowing brightly, surrounded by upward-trending financial charts and graphs.",
+                "A bustling, modern marketplace where digital transactions are visualized as light flowing between people and vendors, symbolizing a digital economy.",
+                "Photorealistic image of a new infrastructure project like a bridge or highway, bathed in golden morning light, representing national growth.",
+            ],
+            'science': [
+                "A stunning view of a satellite orbiting Earth with the Indian subcontinent visible below, against a cosmic background of nebulae and stars.",
+                "A clean, modern laboratory with scientists working with holographic displays showing DNA strands and molecular structures, high-tech aesthetic.",
+                "Digital painting of the Gaganyaan spacecraft soaring through the upper atmosphere, leaving a trail of fire and smoke.",
+                "Microscopic view of new material discovery, with crystalline structures glowing with internal light, abstract and beautiful.",
+            ],
+            'environment': [
+                "Lush green Indian landscape integrating futuristic eco-technology like solar panels that follow the sun and artistic wind turbines.",
+                "A pristine river flowing through a regenerated forest, with clear water and abundant wildlife, symbolizing ecological restoration.",
+                "Concept art of a 'vertical forest' skyscraper in an Indian city, covered in plants and trees, promoting urban greening.",
+                "A split view showing a polluted city on one side and a clean, green city on the other, highlighting the choice for a sustainable future.",
+            ],
+        }
+        prompt_list = prompts.get(category, prompts['technology'])
+        base_prompt = random.choice(prompt_list)
+        return f"{base_prompt}, professional design, high resolution, visually striking for social media."
+
+    def generate_background_image(self, news_data, seed, randomize_seed, width=1024, height=1024, guidance_scale=0.0, num_inference_steps=2):
+        prompt = self.create_image_prompt(news_data)
+        gr.Info(f"🎨 Generating a unique image with prompt: \"{prompt}\"")
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+        generator = torch.Generator().manual_seed(seed)
+        
+        try:
+            image = pipe(
+                prompt=prompt,
+                negative_prompt="",
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                width=width,
+                height=height,
+                generator=generator,
+            ).images[0]
+            return image, seed
+        except Exception as e:
+            gr.Error(f"Image generation failed: {e}. Using fallback image.")
+            return self._create_fallback_image(), seed
+
+    def _create_fallback_image(self):
+        img = Image.new('RGB', (1080, 1080))
+        draw = ImageDraw.Draw(img)
+        for i in range(1080):
+            r, g, b = int(15 + 15 * (i/1080)), int(23 + 18 * (i/1080)), int(42 + 17 * (i/1080))
+            draw.line([(0, i), (1080, i)], fill=(r, g, b))
+        return img
+    
+    def _get_logo(self, size=60):
+        if self.logo_url:
+            try:
+                response = requests.get(self.logo_url)
+                response.raise_for_status()
+                logo_img = Image.open(io.BytesIO(response.content)).convert("RGBA")
+                logo_img.thumbnail((size, size))
+                return logo_img
+            except Exception as e:
+                gr.Warning(f"Failed to load logo from URL: {e}. Using fallback.")
+        logo = Image.new('RGBA', (size, size), (0, 0, 0, 0))
+        draw = ImageDraw.Draw(logo)
+        draw.ellipse([0, 0, size, size], fill=self.brand_colors['accent'])
+        try:
+            font = ImageFont.truetype("seguibld.ttf", int(size * 0.7))
+        except IOError:
+            font = ImageFont.load_default()
+        draw.text((size/2, size/2), "P", font=font, fill='white', anchor="mm")
+        return logo
+    
+    def _draw_text_with_outline(self, draw, position, text, font, fill_color, shadow_color):
+        x, y = position
+        outline = 2
+        for dx in range(-outline, outline + 1):
+            for dy in range(-outline, outline + 1):
+                if dx != 0 or dy != 0:
+                    draw.text((x + dx, y + dy), text, font=font, fill=shadow_color)
+        draw.text(position, text, font=font, fill=fill_color)
+
+    def create_instagram_post(self, category, seed, randomize_seed, width=1024, height=1024, guidance_scale=0.0, num_inference_steps=2):
+        news_data = self.fetch_latest_news(category)
+        if not news_data:
+            gr.Error("Could not retrieve any news story. Please try again.")
+            return None, None, None, seed
+            
+        headline, bullet_points = self.generate_llm_content(news_data)
+        bg_image, seed = self.generate_background_image(news_data, seed, randomize_seed, width, height, guidance_scale, num_inference_steps)
+        bg_image = bg_image.resize((1080, 1080))
+        
+        overlay = Image.new('RGBA', (1080, 1080), (0, 0, 0, 0))
+        draw = ImageDraw.Draw(overlay)
+        draw.rectangle([0, 0, 1080, 1080], fill=(0, 0, 0, 80))
+            
+        try:
+            title_font = ImageFont.truetype("impact.ttf", 85)
+            bullet_font = ImageFont.truetype("verdana.ttf", 36)
+            brand_font = ImageFont.truetype("verdana.ttf", 28)
+            date_font = ImageFont.truetype("verdanab.ttf", 24)
+        except IOError:
+            gr.Warning("Custom fonts not found. Using default fonts.")
+            title_font, bullet_font, brand_font, date_font = (ImageFont.load_default(),)*4
+            
+        draw.text((60, 60), "TODAY'S UPDATE", font=date_font, fill=self.brand_colors['accent'])
+        
+        y_pos = 150
+        for line in textwrap.wrap(headline.upper(), width=25):
+            self._draw_text_with_outline(draw, (60, y_pos), line, title_font, self.brand_colors['text'], self.brand_colors['text_shadow'])
+            y_pos += 95
+            
+        y_pos = 480
+        content_area_end_y = 930
+        stop_drawing = False
+
+        for point_text in bullet_points:
+            if stop_drawing:
+                break
+
+            bullet_radius = 6
+            bullet_x = 70
+            first_line_height = bullet_font.getbbox('A')[3]
+            bullet_y = y_pos + (first_line_height / 2) - bullet_radius
+            
+            if bullet_y > content_area_end_y:
+                break
+
+            draw.ellipse(
+                (bullet_x, bullet_y, bullet_x + bullet_radius*2, bullet_y + bullet_radius*2),
+                fill=self.brand_colors['accent']
+            )
+            
+            text_x = bullet_x + bullet_radius*2 + 20
+            wrapped_text = textwrap.wrap(point_text, width=40)
+            
+            line_y = y_pos
+            for line in wrapped_text:
+                if line_y + first_line_height > content_area_end_y:
+                    stop_drawing = True
+                    break
+                
+                self._draw_text_with_outline(draw, (text_x, line_y), line, bullet_font, self.brand_colors['text'], self.brand_colors['text_shadow'])
+                line_y += 45
+            
+            y_pos = line_y + 20
+
+        draw.rectangle([0, 950, 1080, 1080], fill=self.brand_colors['primary'])
+        logo = self._get_logo(size=100)
+        
+        final_image = Image.alpha_composite(bg_image.convert('RGBA'), overlay)
+        final_image.paste(logo, (940, 965), logo)
+        
+        draw_final = ImageDraw.Draw(final_image)
+        draw_final.text((60, 990), "+91 9354253889", font=brand_font, fill=self.brand_colors['text'], anchor="lm")
+        draw_final.text((920, 990), "pragzo.com", font=brand_font, fill=self.brand_colors['accent'], anchor="rm")
+        
+        return final_image, headline, news_data, seed
+
+# Gradio Interface
+def generate_post(gemini_api_key, news_api_key, logo_url, category, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+    if not gemini_api_key:
+        gr.Error("Please provide a Gemini API key.")
+        return None, None, None, seed
+    
+    generator = CurrentAffairsGenerator(gemini_api_key, news_api_key, logo_url)
+    image, headline, news_data, seed = generator.create_instagram_post(category, seed, randomize_seed, width, height, guidance_scale, num_inference_steps)
+    
+    if image:
+        img_buffer = io.BytesIO()
+        image.save(img_buffer, format='PNG')
+        return image, headline, f"**Original Title:** {news_data['title']}\n**Source:** {news_data['source']}\n**Content Analyzed by AI:**\n\n> {news_data['full_content']}", seed
+    return None, None, None, seed
 
 css = """
-body {
-    font-family: 'Open Sans', sans-serif;
-    background: linear-gradient(135deg, #f5f7fa, #c3cfe2);
-    margin: 0;
-    padding: 0;
-}
-.title {
-    font-size: 1.8em;
-    font-weight: bold;
-    text-align: center;
-    margin: 20px 0;
-}
-footer {
-    visibility: hidden;
+#col-container {
+    margin: 0 auto;
+    max-width: 800px;
 }
 """
 
-@spaces.GPU(duration=120)
-def inference(
-    prompt: str,
-    seed: int,
-    randomize_seed: bool,
-    width: int,
-    height: int,
-    guidance_scale: float,
-    num_inference_steps: int,
-    lora_scale: float,
-    progress: gr.Progress = gr.Progress(track_tqdm=True),
-):
-    try:
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-        generator = torch.Generator(device=device).manual_seed(seed)
-
-        image = pipeline(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
-
-        filepath = save_generated_image(image, prompt)
-        return image, seed, load_generated_images()
-    except Exception as e:
-        print(f"🚨 Generation error: {e}")
-        return Image.new("RGB", (width, height), "white"), seed, load_generated_images()
-
-examples = [
-    "Pepe the frog playing fetch with a golden retriever in a sunny park...",
-    "Pepe the frog dressed in full military gear...",
-    "A medieval Pepe knight in gleaming armor...",
-    "A charismatic Pepe the frog addressing a crowd from a podium...",
-    "Pepe the frog enjoying a peaceful morning at home...",
-    "Businessman Pepe walking confidently through a sleek office lobby...",
-]
-
-with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="emerald"), analytics_enabled=False) as demo:
-    gr.HTML('<div class="title">PEPE Meme Generator</div>')
-
-    gr.HTML("""
-        <a href="https://visitorbadge.io/status?path=https%3A%2F%2Fopenfree-pepe.hf.space">
-            <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fopenfree-pepe.hf.space&countColor=%23263759" />
-        </a>
-    """)
-
-    with gr.Tabs() as tabs:
-        with gr.Tab("Generation"):
-            with gr.Column():
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt",
-                        container=False,
-                    )
-                    run_button = gr.Button("Run", scale=0)
-
-                result = gr.Image(label="Result", show_label=False)
-
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=42,
-                    )
-                    randomize_seed = gr.Checkbox(
-                        label="Randomize seed",
-                        value=True
-                    )
-
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=1024,
-                        )
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=768,
-                        )
-
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=0.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=3.5,
-                        )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=30,
-                        )
-                        lora_scale = gr.Slider(
-                            label="LoRA scale",
-                            minimum=0.0,
-                            maximum=1.0,
-                            step=0.1,
-                            value=1.0,
-                        )
-
-                gr.Examples(
-                    examples=examples,
-                    inputs=[prompt],
-                    outputs=[result, seed],
-                )
-
-        with gr.Tab("Gallery"):
-            gr.Markdown("### Generated Images Gallery")
-            generated_gallery = gr.Gallery(
-                label="Generated Images",
-                columns=6,
-                show_label=False,
-                value=load_generated_images(),
-                elem_id="generated_gallery",
-                height="auto"
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# 🚀 Pragzo | AI Current Affairs Post Generator")
+    gr.Markdown("Instantly create stunning, unique Instagram posts from live news, powered by AI and SDXL-Turbo.")
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            gemini_api_key = gr.Textbox(label="Gemini API Key", type="password", placeholder="Enter your Gemini API key")
+            news_api_key = gr.Textbox(label="News API Key (Optional)", type="password", placeholder="Enter your News API key")
+            logo_url = gr.Textbox(label="Logo URL (Optional)", placeholder="https://.../logo.png")
+
+        with gr.Row():
+            category = gr.Dropdown(
+                label="News Category",
+                choices=["Technology", "Economy", "Science", "Environment", "Sports", "Education"],
+                value="Technology"
             )
-            refresh_btn = gr.Button("🔄 Refresh Gallery")
+            run_button = gr.Button("🛰️ Generate Post", variant="primary")
 
-    def refresh_gallery():
-        return load_generated_images()
+        result = gr.Image(label="Generated Instagram Post", show_label=True)
+        headline = gr.Textbox(label="AI-Generated Headline", interactive=False)
+        news_content = gr.Markdown(label="Source News Content")
 
-    refresh_btn.click(
-        fn=refresh_gallery,
-        inputs=None,
-        outputs=generated_gallery,
-    )
+        with gr.Accordion("Advanced Image Settings", open=False):
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Row():
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+            with gr.Row():
+                guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=0.0)
+                num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=2)
 
     run_button.click(
-        fn=inference,
-        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_scale],
-        outputs=[result, seed, generated_gallery],
-    )
-
-    prompt.submit(
-        fn=inference,
-        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_scale],
-        outputs=[result, seed, generated_gallery],
+        fn=generate_post,
+        inputs=[gemini_api_key, news_api_key, logo_url, category, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, headline, news_content, seed]
     )
 
-demo.queue()
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()