Upload folder using huggingface_hub

app.py

@@ -5,10 +5,20 @@ import os
 from gradio_client import Client
 import asyncio
 import json
+from reportlab.lib.pagesizes import letter
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as RLImage
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.lib.units import inch
+from reportlab.lib import colors
+from diffusers import FluxPipeline
+from PIL import Image
+import io
+import tempfile
 
 # Initialize the text generation pipeline and MCP client
 generator = None
 mcp_client = None
+image_generator = None
 
 # MCP client configuration
 MCP_ENDPOINTS = {
@@ -72,6 +82,34 @@ def initialize_mcp_client():
     except Exception as e:
         return f"MCP client initialization failed: {str(e)}"
 
+def initialize_image_generator():
+    """Initialize FLUX Schnell for image generation"""
+    global image_generator
+    try:
+        # Try to load FLUX Schnell
+        image_generator = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell",
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None
+        )
+        if not torch.cuda.is_available():
+            image_generator = image_generator.to("cpu")
+        return "FLUX Schnell image generator loaded successfully!"
+    except Exception as e:
+        # Fallback to a smaller model if FLUX fails
+        try:
+            from diffusers import StableDiffusionPipeline
+            image_generator = StableDiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                requires_safety_checker=False
+            )
+            image_generator = image_generator.to("cpu")
+            return "Fallback to Stable Diffusion v1.5 for image generation!"
+        except Exception as e2:
+            return f"Image generation initialization failed: {str(e)}, Fallback: {str(e2)}"
+
 def generate_with_mcp(topic, target_audience, key_points, tone, length, model_choice="local"):
     """Generate one-pager using MCP client or local model"""
     
@@ -225,11 +263,130 @@ def import_date():
     from datetime import datetime
     return datetime.now().strftime("%B %d, %Y")
 
+def generate_header_image(topic, tone):
+    """Generate a header image for the one-pager using FLUX Schnell"""
+    global image_generator
+    
+    if image_generator is None:
+        return None
+        
+    try:
+        # Create a professional prompt for the header image
+        image_prompt = f"Professional business illustration for {topic}, {tone.lower()} style, corporate presentation, clean design, business graphics, high quality, no text"
+        
+        # Generate image
+        image = image_generator(
+            prompt=image_prompt,
+            num_inference_steps=4,  # Fast generation with Schnell
+            guidance_scale=0.0,  # FLUX Schnell doesn't need guidance
+            height=256,
+            width=512
+        ).images[0]
+        
+        return image
+        
+    except Exception as e:
+        print(f"Image generation failed: {str(e)}")
+        return None
+
+def export_to_pdf(content, topic, header_image=None):
+    """Export the one-pager content to PDF"""
+    try:
+        # Create a temporary file for the PDF
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
+            pdf_path = tmp_file.name
+        
+        # Create PDF document
+        doc = SimpleDocTemplate(pdf_path, pagesize=letter, topMargin=0.5*inch)
+        styles = getSampleStyleSheet()
+        
+        # Custom styles
+        title_style = ParagraphStyle(
+            'CustomTitle',
+            parent=styles['Heading1'],
+            fontSize=16,
+            spaceAfter=20,
+            textColor=colors.darkblue,
+            alignment=1  # Center alignment
+        )
+        
+        body_style = ParagraphStyle(
+            'CustomBody', 
+            parent=styles['Normal'],
+            fontSize=10,
+            fontName='Courier',  # Monospace font to preserve ASCII formatting
+            leftIndent=0,
+            rightIndent=0
+        )
+        
+        # Build PDF content
+        story = []
+        
+        # Add header image if available
+        if header_image:
+            try:
+                # Save image temporarily
+                img_buffer = io.BytesIO()
+                header_image.save(img_buffer, format='PNG')
+                img_buffer.seek(0)
+                
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as img_file:
+                    img_file.write(img_buffer.getvalue())
+                    img_path = img_file.name
+                
+                # Add image to PDF
+                img = RLImage(img_path, width=6*inch, height=3*inch)
+                story.append(img)
+                story.append(Spacer(1, 20))
+                
+                # Clean up temp image file
+                os.unlink(img_path)
+                
+            except Exception as e:
+                print(f"Failed to add image to PDF: {str(e)}")
+        
+        # Add title
+        story.append(Paragraph(f"Business Document: {topic}", title_style))
+        story.append(Spacer(1, 20))
+        
+        # Add content (preserve formatting)
+        content_lines = content.split('\n')
+        for line in content_lines:
+            if line.strip():
+                story.append(Paragraph(line.replace('<', '&lt;').replace('>', '&gt;'), body_style))
+            else:
+                story.append(Spacer(1, 6))
+        
+        # Build PDF
+        doc.build(story)
+        
+        return pdf_path
+        
+    except Exception as e:
+        print(f"PDF export failed: {str(e)}")
+        return None
+
+def generate_complete_onepager(topic, target_audience, key_points, tone, length, model_choice="local", include_image=True):
+    """Generate complete one-pager with optional image and return both content and PDF"""
+    
+    # Generate the text content
+    content = generate_with_mcp(topic, target_audience, key_points, tone, length, model_choice)
+    
+    # Generate header image if requested
+    header_image = None
+    if include_image and image_generator is not None:
+        header_image = generate_header_image(topic, tone)
+    
+    # Generate PDF
+    pdf_path = export_to_pdf(content, topic, header_image)
+    
+    return content, pdf_path, header_image
+
 # Create the Gradio interface
 def create_interface():
     with gr.Blocks(title="One-Pager Generator", theme=gr.themes.Soft()) as demo:
         gr.Markdown("# 📄 AI One-Pager Generator")
-        gr.Markdown("Generate professional business documents with visual formatting - NOT markdown! Perfect for printing and presentations.")
+        gr.Markdown("Generate professional business documents with FLUX Schnell images and PDF export! Visual formatting, not markdown.")
         
         with gr.Row():
             with gr.Column(scale=1):
@@ -273,16 +430,35 @@ def create_interface():
                     info="Choose between local Qwen model or MCP-connected external services"
                 )
                 
+                include_image_checkbox = gr.Checkbox(
+                    label="Generate Header Image",
+                    value=True,
+                    info="Use FLUX Schnell to generate a professional header image"
+                )
+                
                 generate_btn = gr.Button("🚀 Generate One-Pager", variant="primary")
                 
             with gr.Column(scale=2):
-                output_text = gr.Textbox(
-                    label="Generated One-Pager",
-                    lines=25,
-                    max_lines=35,
-                    show_copy_button=True,
-                    placeholder="Your generated one-pager will appear here..."
-                )
+                with gr.Row():
+                    output_text = gr.Textbox(
+                        label="Generated One-Pager",
+                        lines=20,
+                        max_lines=30,
+                        show_copy_button=True,
+                        placeholder="Your generated one-pager will appear here...",
+                        scale=2
+                    )
+                    generated_image = gr.Image(
+                        label="Header Image",
+                        scale=1,
+                        height=200
+                    )
+                
+                with gr.Row():
+                    pdf_download = gr.File(
+                        label="Download PDF",
+                        visible=False
+                    )
         
         with gr.Row():
             gr.Markdown("""
@@ -295,21 +471,38 @@ def create_interface():
             """)
         
         # Connect the generate button to the function
+        def generate_and_display(topic, audience, keypoints, tone, length, model, include_image):
+            content, pdf_path, header_image = generate_complete_onepager(
+                topic, audience, keypoints, tone, length, model, include_image
+            )
+            
+            # Return outputs for all components
+            pdf_visible = pdf_path is not None
+            return (
+                content,                    # output_text
+                header_image,              # generated_image  
+                gr.File(value=pdf_path, visible=pdf_visible)  # pdf_download
+            )
+        
         generate_btn.click(
-            fn=generate_with_mcp,
-            inputs=[topic_input, audience_input, keypoints_input, tone_dropdown, length_dropdown, model_dropdown],
-            outputs=output_text
+            fn=generate_and_display,
+            inputs=[topic_input, audience_input, keypoints_input, tone_dropdown, length_dropdown, model_dropdown, include_image_checkbox],
+            outputs=[output_text, generated_image, pdf_download]
         )
         
     return demo
 
 # Initialize model and launch
 if __name__ == "__main__":
-    print("🚀 Starting One-Pager Generator with Qwen 2.5-7B and MCP...")
-    print("📥 Loading AI model...")
+    print("🚀 Starting One-Pager Generator with Qwen 2.5-7B, MCP, and FLUX Schnell...")
+    print("📥 Loading AI text model...")
     model_status = initialize_model()
     print(f"✅ {model_status}")
     
+    print("🎨 Loading FLUX Schnell image generator...")
+    image_status = initialize_image_generator()
+    print(f"✅ {image_status}")
+    
     print("🔗 Initializing MCP client...")
     mcp_status = initialize_mcp_client()
     print(f"✅ {mcp_status}")

requirements.txt

@@ -6,4 +6,8 @@ tokenizers
 accelerate
 psutil
 mcp
-gradio_client
+gradio_client
+reportlab
+fpdf2
+diffusers
+pillow