import os import gc import torch import streamlit as st import tempfile import json import subprocess import shutil from datetime import datetime from io import BytesIO import random from PIL import Image import base64 # For decoding images # --- Hugging Face Model Libraries (Local Models for LLM & TTS) --- from transformers import AutoTokenizer, AutoModelForCausalLM from parler_tts import ParlerTTSForConditionalGeneration # --- Google Generative AI (Gemini API) --- try: from google import generativeai as genai from google.generativeai import types as genai_types google_gemini_sdk_available = True except ImportError: google_gemini_sdk_available = False st.error("google-generativeai library not found. Please install it: pip install google-generativeai") # --- Config --- st.set_page_config(layout="wide", page_title="⚡ POV Generator Pro (Gemini SDK Images)") # --- Local Model IDs (LLM & TTS) --- LLM_MODEL_ID = "openai-community/gpt2-medium" TTS_MODEL_ID = "parler-tts/parler-tts-mini-v1.1" # --- Gemini API Configuration (from Streamlit Secrets) --- # Ensure this is set in your Streamlit secrets (e.g., secrets.toml) # Example secrets.toml: # GEMINI_API_KEY = "YOUR_GEMINI_API_KEY" # GEMINI_IMAGE_MODEL_ID = "gemini-2.0-flash-preview-image-generation" # Your custom/preview model GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.environ.get("GEMINI_API_KEY")) # This is the model name you provided. # WARNING: This is not a standard public model name for image generation via this SDK. # It might be a preview model you have access to, or it might not work as expected for image generation. GEMINI_IMAGE_MODEL_ID = st.secrets.get("GEMINI_IMAGE_MODEL_ID", "gemini-2.0-flash-preview-image-generation") # --- Cache and Temp Directory Setup --- CACHE_DIR = os.path.join(tempfile.gettempdir(), "hf_cache_pov_generator") os.makedirs(CACHE_DIR, exist_ok=True) os.environ['HUGGINGFACE_HUB_CACHE'] = CACHE_DIR os.environ['HF_HOME'] = CACHE_DIR os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR # --- Session State Initialization --- if 'run_id' not in st.session_state: st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S") # ... (keep other session state initializations) ... if 'story_data' not in st.session_state: st.session_state.story_data = None if 'pil_images' not in st.session_state: st.session_state.pil_images = None if 'image_paths_for_video' not in st.session_state: st.session_state.image_paths_for_video = None if 'audio_paths' not in st.session_state: st.session_state.audio_paths = None if 'video_path' not in st.session_state: st.session_state.video_path = None if 'temp_base_dir' not in st.session_state: st.session_state.temp_base_dir = None # --- Utility Functions (Keep from previous version) --- def get_session_temp_dir(): if st.session_state.get('temp_base_dir') and os.path.exists(st.session_state.temp_base_dir): return st.session_state.temp_base_dir base_dir = os.path.join(tempfile.gettempdir(), f"pov_generator_run_{st.session_state.run_id}") os.makedirs(base_dir, exist_ok=True) st.session_state.temp_base_dir = base_dir return base_dir def cleanup_temp_files(specific_dir=None): path_to_clean = specific_dir or st.session_state.get("temp_base_dir") if path_to_clean and os.path.exists(path_to_clean): try: shutil.rmtree(path_to_clean) except Exception as e: st.warning(f"Error cleaning temp dir {path_to_clean}: {e}") if specific_dir is None: st.session_state.temp_base_dir = None def clear_torch_cache(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # --- Model Loading (Cached for LLM & TTS - local) --- @st.cache_resource def load_llm_model_and_tokenizer(model_id): # ... (same as before) tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR) model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto", cache_dir=CACHE_DIR ) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token model.config.pad_token_id = model.config.eos_token_id return model, tokenizer @st.cache_resource def load_tts_model_and_tokenizers(model_id): # ... (same as before) tts_model = ParlerTTSForConditionalGeneration.from_pretrained( model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto", cache_dir=CACHE_DIR ) prompt_tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR) desc_tokenizer_path = tts_model.config.text_encoder.name_or_path if hasattr(tts_model.config.text_encoder, 'name_or_path') else tts_model.config.text_encoder._name_or_path desc_tokenizer = AutoTokenizer.from_pretrained(desc_tokenizer_path, cache_dir=CACHE_DIR) return tts_model, prompt_tokenizer, desc_tokenizer # --- Gemini SDK Client (Cached) --- @st.cache_resource def get_gemini_client(api_key): if not google_gemini_sdk_available: return None if not api_key: st.error("GEMINI_API_KEY not found in Streamlit secrets. Please configure it.") return None try: genai.configure(api_key=api_key) # The client itself isn't what we cache, genai.configure sets it globally for the module. # We just need a way to ensure it's configured once. # For actual model interaction, we'll use genai.GenerativeModel return genai # Return the module itself as a signal it's configured except Exception as e: st.error(f"Error configuring Gemini API client: {e}") return None # --- Step 1: Generate JSON Story (Keep from previous version) --- def generate_story(prompt: str, num_scenes: int): # ... (Same as your previous generate_story function) model, tokenizer = load_llm_model_and_tokenizer(LLM_MODEL_ID) story_prompt = ( f"Generate a compelling short POV story based on the following prompt: '{prompt}'. " f"The story should consist of exactly {num_scenes} distinct scenes. " f"Clearly separate each scene with the delimiter '###'. " f"Do not include any introductory or concluding text. Each scene: 2-4 sentences." ) input_ids = tokenizer.encode(story_prompt, return_tensors="pt").to(model.device) max_model_tokens = getattr(model.config, 'n_positions', 1024) max_new = max_model_tokens - input_ids.shape[1] - 20 max_new_tokens_val = min(num_scenes * 75, 700, max_new) if max_new_tokens_val <= 0: st.error("Prompt too long."); return None output = model.generate( input_ids, max_new_tokens=max_new_tokens_val, do_sample=True, temperature=0.7, top_k=50, pad_token_id=tokenizer.eos_token_id ) full_result = tokenizer.decode(output[0], skip_special_tokens=True) generated_text = full_result[len(story_prompt):].strip() if full_result.startswith(story_prompt) else full_result scenes_raw = generated_text.split("###") processed_scenes = [s.strip() for s in scenes_raw if s.strip()] if not processed_scenes: st.error(f"No scenes from LLM: {generated_text}"); return None final_scenes = processed_scenes[:num_scenes] if len(final_scenes) < num_scenes: st.warning(f"LLM made {len(final_scenes)}/{num_scenes} scenes.") clear_torch_cache() return {"title": prompt[:60].capitalize(), "scenes": final_scenes} # --- Step 2: Generate Images with Gemini SDK --- def generate_image_with_gemini_sdk(prompt_text: str, seed: int): # Width/Height might not be controllable if not google_gemini_sdk_available: st.error("Google Generative AI SDK not available.") return None gemini_configured_module = get_gemini_client(GEMINI_API_KEY) if not gemini_configured_module: return None try: # IMPORTANT: The model name GEMINI_IMAGE_MODEL_ID is user-provided and may not be a standard image generation model. # The `response_modalities` config is also based on the user's snippet. # This whole block assumes this model and config work as intended by the user. model_instance = genai.GenerativeModel(GEMINI_IMAGE_MODEL_ID) st.write(f"Using Gemini Model for Image: {GEMINI_IMAGE_MODEL_ID}") # For debugging # The concept of 'seed' might not be directly applicable or exposed in the same way # for `generate_content` with some Gemini models. We include it if the API might use it. # GenerationConfig can also take 'temperature', 'top_p', 'top_k', 'candidate_count', 'max_output_tokens'. # 'seed' is not a standard param for GenerateContentConfig in the public SDK. # We will use a basic config based on your snippet. generation_config = genai_types.GenerationConfig( # response_mime_type="image/png", # Often used with specific image models, but response_modalities is different # Not adding seed here as it's not standard for this config object. # Temperature, etc., could be added if relevant for image quality. ) # The `contents` structure might need to be more complex if the model expects specific formatting. # For a simple prompt, a string might suffice. # The `response_modalities` part is from your snippet, which is unusual for GenerateContentConfig. # It's typically set at the model level or inferred. # Let's try to construct the request as closely as possible to your example structure. # `client.models.generate_content` is not the typical way; it's usually `GenerativeModel_instance.generate_content` # Corrected way to use generate_content with a GenerativeModel instance: response = model_instance.generate_content( contents=[prompt_text], # `contents` should be an iterable of Content objects or strings generation_config=generation_config, # The `response_modalities` is not a direct parameter of `generate_content`. # It's usually an expectation from the model. If your model *requires* it, # it might need to be part of a more complex `Content` object within `contents`. # For now, we'll assume the model inherently knows to produce an image if prompted correctly, # or that `GEMINI_IMAGE_MODEL_ID` implies image output. ) # Process the response to find image data # A response can have multiple `parts`. Image data is typically in a part with a specific mime_type. if response.parts: for part in response.parts: if part.mime_type and part.mime_type.startswith("image/"): # Image data is in `part.data` (bytes) or `part.inline_data.data` image_bytes = None if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): # Common structure image_bytes = part.inline_data.data elif hasattr(part, 'data'): # Fallback image_bytes = part.data if image_bytes: pil_image = Image.open(BytesIO(image_bytes)) return pil_image else: st.warning("Image part found in Gemini response, but no image data.") st.error("No image data found in Gemini response parts, though parts existed.") st.json([part_to_dict(p) for p in response.parts]) # Show what was received return None else: st.error("Gemini SDK returned no parts in the response.") if response.prompt_feedback: st.warning(f"Prompt Feedback: {response.prompt_feedback}") return None except Exception as e: st.error(f"Error during Gemini SDK image generation: {e}") st.error(f"Model used: {GEMINI_IMAGE_MODEL_ID}. Prompt: {prompt_text[:100]}...") if "API key not valid" in str(e) or "PERMISSION_DENIED" in str(e): st.error("Please check your GEMINI_API_KEY and ensure it's valid and has permissions for this model.") elif "Could not find model" in str(e) or "MODEL_NAME_INVALID" in str(e): st.error(f"The model '{GEMINI_IMAGE_MODEL_ID}' could not be found or is invalid for this API. Please verify the model name.") elif "DeadlineExceeded" in str(e): st.error("The request to Gemini API timed out. The model might be slow or the task too complex.") # You might want to inspect the full error for more details # print(f"Full Gemini API Error: {e}") return None def part_to_dict(part): """Helper to convert a Part object to a dict for st.json, handling bytes.""" data = " данни" # Placeholder if data is binary and large if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): data = f"" if part.inline_data.data else "няма данни" elif hasattr(part, 'data'): data = f"" if part.data else "няма данни" return { "text": part.text if hasattr(part, 'text') else None, "mime_type": part.mime_type if hasattr(part, 'mime_type') else None, "data_summary": data } def generate_images_for_scenes(scenes, randomize_seed_img, base_seed_img): # img_width, img_height removed as likely not controllable pil_images = [] frames_dir = os.path.join(get_session_temp_dir(), "frames_for_video") os.makedirs(frames_dir, exist_ok=True) image_paths_for_video = [] cols = st.columns(min(3, len(scenes))) col_idx = 0 st.warning(f""" **Note on Image Generation with `{GEMINI_IMAGE_MODEL_ID}`:** - The model name provided is not a standard public image generation model via this SDK. Results may vary. - Image dimensions (width/height) are determined by the model. UI sliders for width/height are not used for this generation method. - The 'seed' parameter's effect depends on the specific model's implementation via this API. """) for i, scene_text in enumerate(scenes): # Seed concept might be different or not directly exposed for `generate_content` current_seed = random.randint(0, 9_999_999) if randomize_seed_img else base_seed_img + i with st.spinner(f"Generating image for scene {i+1} via Gemini SDK (seed hint: {current_seed})..."): img = None try: # The prompt needs to clearly ask for an image based on the scene text. image_prompt = f"Visually represent this scene: {scene_text}. Cinematic, detailed, high quality." img = generate_image_with_gemini_sdk(image_prompt, current_seed) except Exception as e: st.error(f"Failed to generate image for scene {i+1}: {e}") img = None pil_images.append(img) if img: st.caption(f"Generated image size: {img.width}x{img.height}") img_path = os.path.join(frames_dir, f"frame_{i:03d}.png") img.save(img_path) image_paths_for_video.append(img_path) with cols[col_idx % len(cols)]: st.image(img, caption=f"Scene {i+1} (Seed hint: {current_seed})") img_byte_arr = BytesIO() img.save(img_byte_arr, format='PNG') st.download_button( label=f"Download Img {i+1}", data=img_byte_arr.getvalue(), file_name=f"scene_{i+1}_image.png", mime="image/png", key=f"dl_img_{i}" ) col_idx += 1 else: image_paths_for_video.append(None) with cols[col_idx % len(cols)]: st.warning(f"Image for Scene {i+1} could not be generated.") col_idx += 1 return pil_images, image_paths_for_video # --- Step 3: Generate TTS (Keep from previous version) --- def generate_audios_for_scenes(scenes): # ... (Same as your previous generate_audios_for_scenes function) tts_model, prompt_tokenizer, desc_tokenizer = load_tts_model_and_tokenizers(TTS_MODEL_ID) audio_dir = os.path.join(get_session_temp_dir(), "audio_files") os.makedirs(audio_dir, exist_ok=True) audio_paths = [] cols = st.columns(min(3, len(scenes))) col_idx = 0 tts_description = "A neutral and clear narrator voice." for i, scene_text in enumerate(scenes): with st.spinner(f"Generating audio for scene {i+1}..."): try: desc_ids = desc_tokenizer(tts_description, return_tensors="pt").input_ids.to(tts_model.device) prompt_ids = prompt_tokenizer(scene_text, return_tensors="pt").input_ids.to(tts_model.device) generation_output = tts_model.generate(input_ids=desc_ids, prompt_input_ids=prompt_ids) audio_waveform = generation_output.to(torch.float32).cpu().numpy() file_path = os.path.join(audio_dir, f"audio_scene_{i+1}.wav") sf.write(file_path, audio_waveform, tts_model.config.sampling_rate) audio_paths.append(file_path) with cols[col_idx % len(cols)]: st.markdown(f"**Audio Scene {i+1}**"); st.audio(file_path) with open(file_path, "rb") as f_audio: st.download_button( f"Download Audio {i+1}", f_audio.read(), f"s_{i+1}_audio.wav", "audio/wav", key=f"dl_aud_{i}" ) col_idx +=1 except Exception as e: st.error(f"Audio error s{i+1}: {e}"); audio_paths.append(None) clear_torch_cache() return audio_paths # --- Step 4: Create Video (Keep from previous version) --- def create_video_from_scenes(image_file_paths, audio_file_paths, output_filename="final_pov_video.mp4"): # ... (Same as your previous create_video_from_scenes function) if not image_file_paths or not audio_file_paths or len(image_file_paths) != len(audio_file_paths): st.error("Mismatch/missing assets for video."); return None try: subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True) except: st.error("FFMPEG not found."); return None temp_clips_dir = os.path.join(get_session_temp_dir(), "temp_video_clips") os.makedirs(temp_clips_dir, exist_ok=True) clips_paths, valid_scenes = [], 0 for i, (img_p, aud_p) in enumerate(zip(image_file_paths, audio_file_paths)): if not (img_p and aud_p): continue try: aud_info = sf.info(aud_p); aud_dur = aud_info.duration if aud_dur <= 0.1: aud_dur = 1.0 clip_p = os.path.join(temp_clips_dir, f"c_{i:03d}.mp4") cmd = [ "ffmpeg", "-y", "-loop", "1", "-i", img_p, "-i", aud_p, "-c:v", "libx264", "-preset", "medium", "-tune", "stillimage", "-c:a", "aac", "-b:a", "192k", "-pix_fmt", "yuv420p", "-t", str(aud_dur), "-shortest", clip_p ] res = subprocess.run(cmd, capture_output=True, text=True) if res.returncode != 0: st.error(f"FFMPEG clip {i+1} err:\n{res.stderr}"); continue clips_paths.append(clip_p); valid_scenes +=1 except Exception as e: st.error(f"Video scene {i+1} err: {e}") if not clips_paths or valid_scenes == 0: st.error("No valid video clips."); return None concat_list_f = os.path.join(temp_clips_dir, "concat_list.txt") with open(concat_list_f, "w") as f: for cp in clips_paths: f.write(f"file '{os.path.basename(cp)}'\n") final_vid_p = os.path.join(get_session_temp_dir(), output_filename) concat_cmd = ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_list_f, "-c", "copy", final_vid_p] res = subprocess.run(concat_cmd, capture_output=True, text=True, cwd=temp_clips_dir) if res.returncode != 0: st.error(f"FFMPEG concat err:\n{res.stderr}"); return None st.success("Video created!"); return final_vid_p # --- Main App UI --- st.title("⚡ POV Story Generator Pro (Gemini SDK Images)") st.markdown("Create POV stories: Local LLM/TTS & Google Gemini SDK for Image Generation.") if not google_gemini_sdk_available: st.error("Google Generative AI SDK (`google-generativeai`) is not installed. Image generation will not be available.") st.markdown("Please install it: `pip install google-generativeai`") if not GEMINI_API_KEY: st.warning("`GEMINI_API_KEY` not set in secrets. Please configure it for Gemini SDK image generation.") with st.sidebar: st.header("📝 Story Configuration") prompt = st.text_area( "POV story prompt:", st.session_state.get("user_prompt", "A squirrel discovers a tiny spaceship."), height=100, key="user_prompt_input" ) num_scenes = st.slider("Number of Scenes:", 1, 8, st.session_state.get("num_scenes_val", 2), key="num_scenes_slider") st.header("🎨 Image Generation (Gemini SDK)") st.caption(f"Using Gemini Model: `{GEMINI_IMAGE_MODEL_ID}`") st.warning("Note: The specified Gemini model for images may be experimental. Image dimensions (width/height) are determined by the model.") # Width/Height sliders are kept for conceptual consistency but noted as not directly controlling this Gemini API. img_width_ui = st.slider("Image Width (Informational):", 512, 1536, st.session_state.get("img_width_val", 1024), 128, key="img_w_ui") img_height_ui = st.slider("Image Height (Informational):", 512, 1536, st.session_state.get("img_height_val", 1024), 128, key="img_h_ui") cols_seed = st.columns(2) with cols_seed[0]: base_seed_img = st.number_input("Base Image Seed (hint):", 0, 9999999, st.session_state.get("base_seed_val", 12345), key="base_seed_in") with cols_seed[1]: randomize_seed_img = st.checkbox("Randomize Seed (hint)", st.session_state.get("random_seed_bool", True), key="rand_seed_chk") st.markdown("---") can_generate = google_gemini_sdk_available and GEMINI_API_KEY if st.button("🚀 Generate Full Story & Assets", type="primary", use_container_width=True, disabled=not can_generate): st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S") cleanup_temp_files() st.session_state.story_data = None; st.session_state.pil_images = None st.session_state.image_paths_for_video = None; st.session_state.audio_paths = None st.session_state.video_path = None st.session_state.user_prompt = prompt; st.session_state.num_scenes_val = num_scenes # Store informational width/height, though not used directly in Gemini SDK call here st.session_state.img_width_val = img_width_ui; st.session_state.img_height_val = img_height_ui st.session_state.base_seed_val = base_seed_img; st.session_state.random_seed_bool = randomize_seed_img st.session_state.generate_all = True elif not can_generate: st.warning("Image generation disabled: missing Gemini API Key or SDK.") st.markdown("---") st.header("🛠️ Utilities") if st.button("🧹 Clear Cache & Temp Files & Restart", use_container_width=True): st.cache_resource.clear() keys_to_clear = ['story_data', 'pil_images', 'image_paths_for_video', 'audio_paths', 'video_path', 'temp_base_dir', 'generate_all'] for key in keys_to_clear: if key in st.session_state: del st.session_state[key] cleanup_temp_files(); st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S") st.success("Caches & temp files cleared. App restarting."); st.rerun() # Main content area - (generation logic similar to previous, adapted for new image function) if st.session_state.get("generate_all"): # --- 1. Story --- with st.status("🧠 Generating story...", True) as s_story: try: st.session_state.story_data = generate_story(st.session_state.user_prompt, st.session_state.num_scenes_val) if st.session_state.story_data: s_story.update(label="Story generated!", state="complete") else: s_story.update(label="Story fail.", state="error"); st.session_state.generate_all = False except Exception as e: st.error(f"Story err: {e}"); s_story.update(label="Story err.",state="error"); st.session_state.generate_all=False if st.session_state.story_data: st.subheader(f"🎬 Story: {st.session_state.story_data['title']}") for i, scene_text in enumerate(st.session_state.story_data['scenes']): st.markdown(f"**Scene {i+1}:** {scene_text}") story_json = json.dumps(st.session_state.story_data, indent=2) st.download_button("Download Story (JSON)", story_json, f"story.json", "application/json") st.markdown("---") # --- 2. Images --- if st.session_state.get("generate_all") and st.session_state.story_data: with st.status(f"🎨 Generating images via Gemini SDK ({GEMINI_IMAGE_MODEL_ID})...", True) as s_images: try: st.session_state.pil_images, st.session_state.image_paths_for_video = generate_images_for_scenes( st.session_state.story_data['scenes'], st.session_state.random_seed_bool, st.session_state.base_seed_val ) if st.session_state.pil_images and any(img for img in st.session_state.pil_images): s_images.update(label="Images generated!", state="complete") else: s_images.update(label="Image gen failed/no images.", state="error"); st.session_state.generate_all = False except Exception as e: st.error(f"Image gen err: {e}"); s_images.update(label="Image error.",state="error"); st.session_state.generate_all=False st.markdown("---") # --- 3. Audio --- if st.session_state.get("generate_all") and st.session_state.story_data: with st.status("🔊 Generating audio...", True) as s_audio: try: st.session_state.audio_paths = generate_audios_for_scenes(st.session_state.story_data['scenes']) if st.session_state.audio_paths and any(p for p in st.session_state.audio_paths): s_audio.update(label="Audio generated!", state="complete") else: s_audio.update(label="Audio gen failed.", state="error"); st.session_state.generate_all = False except Exception as e: st.error(f"Audio err: {e}"); s_audio.update(label="Audio error.",state="error"); st.session_state.generate_all=False st.markdown("---") # --- 4. Video --- if st.session_state.get("generate_all") and st.session_state.image_paths_for_video and st.session_state.audio_paths: valid_assets = sum(1 for im,au in zip(st.session_state.image_paths_for_video, st.session_state.audio_paths) if im and au) if valid_assets > 0: with st.status("📹 Creating video...", True) as s_video: try: st.session_state.video_path = create_video_from_scenes( st.session_state.image_paths_for_video, st.session_state.audio_paths ) if st.session_state.video__path: s_video.update(label="Video created!", state="complete") else: s_video.update(label="Video creation failed.", state="error") except Exception as e: st.error(f"Video err: {e}"); s_video.update(label="Video error.", state="error") if st.session_state.video_path: st.subheader("🎞️ Final Video"); st.video(st.session_state.video_path) with open(st.session_state.video_path, "rb") as fv: st.download_button("Download Video", fv.read(), os.path.basename(st.session_state.video_path), "video/mp4") st.markdown("---") else: st.warning("Not enough assets for video.") if "generate_all" in st.session_state: del st.session_state.generate_all elif not st.session_state.get("user_prompt"): st.info("Configure story in sidebar & click 'Generate' to begin!")