import gradio as gr from transformers import pipeline, AutoProcessor, AutoModelForCausalLM from diffusers import StableDiffusionPipeline, DiffusionPipeline import torch from PIL import Image import numpy as np import os import tempfile import moviepy.editor as mpe import nltk from pydub import AudioSegment import warnings import asyncio import edge_tts import random from datetime import datetime import pytz import re import json from gradio_client import Client warnings.filterwarnings("ignore", category=UserWarning) # Ensure NLTK data is downloaded nltk.download('punkt') # Initialize clients arxiv_client = None def init_arxiv_client(): global arxiv_client if arxiv_client is None: arxiv_client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern") return arxiv_client # File I/O Functions def generate_filename(prompt, timestamp=None): """Generate a safe filename from prompt and timestamp""" if timestamp is None: timestamp = datetime.now(pytz.UTC).strftime("%Y%m%d_%H%M%S") # Clean the prompt to create a safe filename safe_prompt = re.sub(r'[^\w\s-]', '', prompt)[:50].strip() return f"story_{timestamp}_{safe_prompt}.txt" def save_story(story, prompt, filename=None): """Save story to file with metadata""" if filename is None: filename = generate_filename(prompt) try: with open(filename, 'w', encoding='utf-8') as f: metadata = { 'timestamp': datetime.now().isoformat(), 'prompt': prompt, 'type': 'story' } f.write(json.dumps(metadata) + '\n---\n' + story) return filename except Exception as e: print(f"Error saving story: {e}") return None def load_story(filename): """Load story and metadata from file""" try: with open(filename, 'r', encoding='utf-8') as f: content = f.read() parts = content.split('\n---\n') if len(parts) == 2: metadata = json.loads(parts[0]) story = parts[1] return metadata, story return None, content except Exception as e: print(f"Error loading story: {e}") return None, None # Story Generation Functions def generate_story(prompt, model_choice): """Generate story using specified model""" try: client = init_arxiv_client() if client is None: return "Error: Story generation service is not available." result = client.predict( prompt=prompt, llm_model_picked=model_choice, stream_outputs=True, api_name="/ask_llm" ) return result except Exception as e: return f"Error generating story: {str(e)}" async def generate_speech(text, voice="en-US-AriaNeural"): """Generate speech from text""" try: communicate = edge_tts.Communicate(text, voice) with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file: tmp_path = tmp_file.name await communicate.save(tmp_path) return tmp_path except Exception as e: print(f"Error in text2speech: {str(e)}") return None def process_story_and_audio(prompt, model_choice): """Process story and generate audio""" try: # Generate story story = generate_story(prompt, model_choice) if isinstance(story, str) and story.startswith("Error"): return story, None, None # Save story filename = save_story(story, prompt) # Generate audio audio_path = asyncio.run(generate_speech(story)) return story, audio_path, filename except Exception as e: return f"Error: {str(e)}", None, None # Main App Code (your existing code remains here) # LLM Inference Class and other existing classes remain unchanged class LLMInferenceNode: # Your existing LLMInferenceNode implementation pass # Initialize models (your existing initialization code remains here) device = "cuda" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if device == "cuda" else torch.float32 # Story generator story_generator = pipeline( 'text-generation', model='gpt2-large', device=0 if device == 'cuda' else -1 ) # Stable Diffusion model sd_pipe = StableDiffusionPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", torch_dtype=torch_dtype ).to(device) # Create the enhanced Gradio interface with gr.Blocks() as demo: gr.Markdown("""# 🎨 AI Creative Suite Generate videos, stories, and more with AI! """) with gr.Tabs(): # Your existing video generation tab with gr.Tab("Video Generation"): with gr.Row(): with gr.Column(): prompt_input = gr.Textbox(label="Enter a Prompt", lines=2) generate_button = gr.Button("Generate Video") with gr.Column(): video_output = gr.Video(label="Generated Video") generate_button.click(fn=process_pipeline, inputs=prompt_input, outputs=video_output) # New story generation tab with gr.Tab("Story Generation"): with gr.Row(): with gr.Column(): story_prompt = gr.Textbox( label="Story Concept", placeholder="Enter your story idea...", lines=3 ) model_choice = gr.Dropdown( label="Model", choices=[ "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.2" ], value="mistralai/Mixtral-8x7B-Instruct-v0.1" ) generate_story_btn = gr.Button("Generate Story") with gr.Row(): story_output = gr.Textbox( label="Generated Story", lines=10, interactive=False ) with gr.Row(): audio_output = gr.Audio( label="Story Narration", type="filepath" ) filename_output = gr.Textbox( label="Saved Filename", interactive=False ) generate_story_btn.click( fn=process_story_and_audio, inputs=[story_prompt, model_choice], outputs=[story_output, audio_output, filename_output] ) # File management section with gr.Row(): file_list = gr.Dropdown( label="Saved Stories", choices=[f for f in os.listdir() if f.startswith("story_") and f.endswith(".txt")], interactive=True ) refresh_btn = gr.Button("🔄 Refresh") def refresh_files(): return gr.Dropdown(choices=[f for f in os.listdir() if f.startswith("story_") and f.endswith(".txt")]) refresh_btn.click(fn=refresh_files, outputs=[file_list]) # Launch the app if __name__ == "__main__": demo.launch(debug=True)