app.py

# -*- coding: utf-8 -*-

import gradio as gr
from huggingface_hub import InferenceClient
from gradio_client import Client
import os
import requests
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor

# Logging configuration
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

# API configuration
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN"))
IMAGE_API_URL = "http://211.233.58.201:7896"

def generate_image_prompt(text: str) -> str:
    """Generate image prompt from novel content"""
    try:
        prompt_messages = [
            {"role": "system", "content": "Extract the most visually descriptive scene or key elements from the given martial arts fantasy novel text and create a detailed image generation prompt."},
            {"role": "user", "content": f"Create an image generation prompt from this text: {text}"}
        ]
        
        response = hf_client.chat_completion(prompt_messages, max_tokens=200)
        image_prompt = response.choices[0].message.content
        return f"traditional chinese martial arts fantasy style, {image_prompt}"
    except Exception as e:
        logging.error(f"Image prompt generation failed: {str(e)}")
        return f"traditional chinese martial arts fantasy style, {text[:200]}"

def generate_image(prompt: str) -> tuple:
    """Image generation function"""
    try:
        client = Client(IMAGE_API_URL)
        result = client.predict(
            prompt=prompt,
            width=768,
            height=768,
            guidance=7.5,
            inference_steps=30,
            seed=3,
            do_img2img=False,
            init_image=None,
            image2image_strength=0.8,
            resize_img=True,
            api_name="/generate_image"
        )
        return result[0], result[1]
    except Exception as e:
        logging.error(f"Image generation failed: {str(e)}")
        return None, f"Error: {str(e)}"

# Global list to store image history
image_history = []

def format_text(text: str, max_line_length: int = 80) -> str:
    """Text formatting function"""
    lines = []
    current_line = ""
    
    for paragraph in text.split('\n'):
        words = paragraph.split()
        
        for word in words:
            if len(current_line) + len(word) + 1 <= max_line_length:
                current_line += word + " "
            else:
                lines.append(current_line.strip())
                current_line = word + " "
        
        if current_line:
            lines.append(current_line.strip())
            current_line = ""
        
        lines.append("")  # Empty line for paragraph separation
    
    return "\n".join(lines)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message="",  
    max_tokens=7860,  
    temperature=0.8, 
    top_p=0.9,  
):
    global image_history
    
    system_prefix = """
You are Martial Arts AI⚔️, a specialized AI focused on creating immersive martial arts fantasy novels. Your responses should start with 'Martial Arts AI⚔️:' and maintain the rich traditions of martial arts fiction while being engaging and creative.

Essential Guidelines:
1. Create continuous, engaging martial arts fantasy narratives up to 7860 tokens
2. Each response should naturally connect with previous content
3. Include these key elements in every response:
   - Detailed martial arts system descriptions
   - Character's internal energy cultivation methods
   - Martial world rules and ethics
   - Dynamic combat scene descriptions
   - Engaging plot development
   - Balance between dialogue and narration

Core Martial Arts Elements:
- Combat Systems (Internal Energy, External Techniques, Lightness Skills, Hidden Weapons)
- Martial Sects (Orthodox, Unorthodox, Neutral)
- Secret Manuals (Ultimate Techniques, Forbidden Arts)
- Power Groups (Sects, Clans, Schools, Families)
- Ancient Martial Arts Secrets and Legends
- Warriors and Masters
- Martial World Leaders

Narrative Style:
1. Clear paragraph breaks with appropriate spacing
2. Dynamic dialogue with character movements
3. Detailed combat choreography
4. Internal energy circulation descriptions
5. Rich environmental descriptions

References Include:
- Martial Arts Techniques
- Martial World Proverbs
- Secret Manual Passages
- Sect Principles
- Martial World Rules
- Combat Formulas
- Sect Documents

[Detailed sect and school information follows in structured format...]

Story Structure:
1. Opening: Fateful martial arts encounter
2. Development: Training and emerging conflicts
3. Crisis: Life-or-death confrontation
4. Climax: Achieving martial arts mastery
5. Resolution: Birth of a new legend

Technical Terms:
- Internal Energy: Meridians, Dantian, Energy Channels
- External Techniques: Fist Methods, Palm Techniques, Finger Skills
- Movement Arts: Lightness Skills, Body Methods
- Hidden Weapons: Concealment, Deployment Methods
- Cultivation Levels: 
  * Minor Achievement, Major Achievement, Transcendent State
  * Mystery Gate, Mystery Portal, Mystery Body
  * Life Gate, Life Portal, Life Body
  * Spirit Gate, Spirit Portal, Spirit Body

Each response should be complete like a chapter while leaving room for continuation."""

    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    current_response = ""
    new_history = history.copy()

    try:
        for msg in hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = msg.choices[0].delta.content
            if token is not None:
                current_response += token
                formatted_response = format_text(current_response)
                new_history = history + [(message, formatted_response)]
                yield new_history, None, [img[0] for img in image_history]

        final_response = format_text(current_response)
        new_history = history + [(message, final_response)]
        
        image_prompt = generate_image_prompt(current_response)
        image, _ = generate_image(image_prompt)
        
        if image is not None:
            image_history.append((image, image_prompt))
        
        yield new_history, image, [img[0] for img in image_history]

    except Exception as e:
        error_message = f"Error: {str(e)}"
        yield history + [(message, error_message)], None, [img[0] for img in image_history]

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css="""
    .message-wrap {
        font-size: 14px !important;
        line-height: 1.5em !important;
        max-width: 90% !important;
        margin: 0 auto !important;
    }
    .message {
        padding: 1em !important;
        margin-bottom: 0.5em !important;
        white-space: pre-wrap !important;
        word-wrap: break-word !important;
        max-width: 100% !important;
    }
    .message p {
        margin: 0 !important;
        padding: 0 !important;
        width: 100% !important;
    }
    .chatbot {
        font-family: 'Noto Sans', sans-serif !important;
    }
    footer {
        visibility: hidden;
    }
        
""") as interface:
    gr.Markdown("# Kungfu Graphic Novel Generator")
    gr.Markdown("### After each chapter is generated, corresponding images are created automatically. Click 'Continue Story' to proceed with the narrative.")    

    gr.HTML("""
        <a href="https://visitorbadge.io/status?path=https%3A%2F%2Ffantaxy-novel-sorim-en.hf.space">
            <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Ffantaxy-novel-sorim-en.hf.space&countColor=%23263759" />
        </a>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                value=[],
                show_label=True,
                label="Story Progress",
                height=500,
                elem_classes="chatbot"
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Enter your prompt",
                    placeholder="Type your story prompt here...",
                    lines=2
                )
                submit_btn = gr.Button("Generate", variant="primary")
            
            system_msg = gr.Textbox(
                label="System Message",
                value="Generate engaging martial arts fantasy content.",
                lines=2
            )
            
            with gr.Row():
                max_tokens = gr.Slider(
                    minimum=1,
                    maximum=8000,
                    value=7000,
                    label="Story Length (tokens)"
                )
                temperature = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.7,
                    label="Creativity Level"
                )
                top_p = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.9,
                    label="Response Focus"
                )
        
        with gr.Column(scale=1):
            image_output = gr.Image(
                label="Latest Scene Illustration",
                height=400
            )
            gallery = gr.Gallery(
                label="Story Illustrations Gallery",
                show_label=True,
                elem_id="gallery",
                columns=[2],
                rows=[2],
                height=300
            )

    examples = gr.Examples(
        examples=[
            ["Continue the story"],
            ["Suggest 10 interesting plot elements for the story"],
        ],
        inputs=msg
    )

    submit_btn.click(
        fn=respond,
        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p],
        outputs=[chatbot, image_output, gallery]
    )
    
    msg.submit(
        fn=respond,
        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p],
        outputs=[chatbot, image_output, gallery]
    )

if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )