import gradio as gr import torch from transformers import pipeline, set_seed from diffusers import StableDiffusionPipeline import os import time # ---- 配置与模型加载 (在应用启动时加载一次) ---- # 检查是否有可用的GPU，否则使用CPU device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") # 1. 语音转文本模型 (Whisper) - 加分项 asr_pipeline = None try: print("Loading ASR pipeline (Whisper)...") # 使用较小的模型以节省资源，可根据需要替换 openai/whisper-medium 或 large # 在不需要GPU的应用部分可以强制使用CPU asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device if device == "cuda" else -1) # whisper在CPU上也可以运行 print("ASR pipeline loaded.") except Exception as e: print(f"Could not load ASR pipeline: {e}. Voice input will be disabled.") # 2. 提示词增强模型 (LLM) - Step 1 prompt_enhancer_pipeline = None try: print("Loading Prompt Enhancer pipeline (GPT-2)...") # 使用 GPT-2 作为示例，实际应用中建议使用更强大的指令微调模型如 Mistral 或 Llama # 注意：GPT-2 可能不会生成特别高质量的SD提示词，这里仅作结构演示 # 如果资源允许，可以替换为 'mistralai/Mistral-7B-Instruct-v0.1' 等，但需要更多内存/GPU prompt_enhancer_pipeline = pipeline("text-generation", model="gpt2", device=device if device == "cuda" else -1) # text-generation在CPU上也可以运行 print("Prompt Enhancer pipeline loaded.") except Exception as e: print(f"Could not load Prompt Enhancer pipeline: {e}. Prompt enhancement might fail.") # 3. 文本到图像模型 (Stable Diffusion) - Step 2 image_generator_pipe = None try: print("Loading Stable Diffusion pipeline (v1.5)...") model_id = "runwayml/stable-diffusion-v1-5" image_generator_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32) image_generator_pipe = image_generator_pipe.to(device) # 如果内存不足，可以启用CPU offloading (需要 accelerate库) # image_generator_pipe.enable_model_cpu_offload() print("Stable Diffusion pipeline loaded.") except Exception as e: print(f"Could not load Stable Diffusion pipeline: {e}. Image generation will fail.") # 如果模型加载失败，创建一个虚拟对象以避免后续代码出错 class DummyPipe: def __call__(self, *args, **kwargs): # 返回一个占位符错误信息或图像 raise RuntimeError(f"Stable Diffusion model failed to load: {e}") image_generator_pipe = DummyPipe() # ---- 核心功能函数 ---- # Step 1: Prompt-to-Prompt def enhance_prompt(short_prompt, style_modifier="cinematic", quality_boost="photorealistic, highly detailed"): """使用LLM增强简短描述""" if not prompt_enhancer_pipeline: return f"[Error: LLM not loaded] Original prompt: {short_prompt}" if not short_prompt: return "[Error: Input description is empty]" # 构建给LLM的指令 # 注意：这个指令对GPT-2来说可能太复杂，对Mistral等更有效 input_text = ( f"Generate a detailed and vivid prompt for an AI image generator based on the following description. " f"Incorporate the style '{style_modifier}' and quality boost '{quality_boost}'. " f"Focus on visual details, lighting, composition, and mood. " f"Description: \"{short_prompt}\"\n\n" f"Detailed Prompt:" ) try: # 设置种子以获得可复现的（某种程度上的）结果 set_seed(int(time.time())) # max_length 控制生成文本的总长度 (包括输入) # num_return_sequences 返回多少个结果 # temperature 控制随机性，较低的值更保守 # no_repeat_ngram_size 避免重复短语 outputs = prompt_enhancer_pipeline( input_text, max_length=150, # 限制输出长度，避免过长 num_return_sequences=1, temperature=0.7, no_repeat_ngram_size=2, pad_token_id=prompt_enhancer_pipeline.tokenizer.eos_token_id # 避免padding warning ) generated_text = outputs[0]['generated_text'] # 从LLM的完整输出中提取增强后的提示词部分 # 简单方法：取 "Detailed Prompt:" 之后的内容 enhanced = generated_text.split("Detailed Prompt:")[-1].strip() # 进一步清理可能包含的原始输入或指令痕迹 if short_prompt in enhanced[:len(short_prompt)+5]: # 如果开头包含原始输入 enhanced = enhanced.replace(short_prompt, "", 1).strip(' ,"') # 添加基础的风格和质量词，如果LLM没有包含的话 if style_modifier not in enhanced: enhanced += f", {style_modifier}" if quality_boost not in enhanced: enhanced += f", {quality_boost}" return enhanced except Exception as e: print(f"Error during prompt enhancement: {e}") return f"[Error: Prompt enhancement failed] Original prompt: {short_prompt}" # Step 2: Prompt-to-Image def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps): """使用Stable Diffusion生成图像""" if not isinstance(image_generator_pipe, StableDiffusionPipeline): raise gr.Error(f"Stable Diffusion model is not available. Load error: {image_generator_pipe}") # 使用gr.Error在UI上显示错误 if not prompt or "[Error:" in prompt: raise gr.Error("Cannot generate image due to invalid or missing prompt.") print(f"Generating image for prompt: {prompt}") print(f"Negative prompt: {negative_prompt}") print(f"Guidance scale: {guidance_scale}, Steps: {num_inference_steps}") try: # 设置随机种子 generator = torch.Generator(device=device).manual_seed(int(time.time())) # 执行推理 with torch.inference_mode(): # 节省内存 image = image_generator_pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=float(guidance_scale), num_inference_steps=int(num_inference_steps), generator=generator ).images[0] print("Image generated successfully.") return image except Exception as e: print(f"Error during image generation: {e}") # 将底层错误传递给 Gradio，使其能在 UI 中显示 raise gr.Error(f"Image generation failed: {e}") # Bonus: Voice-to-Text def transcribe_audio(audio_file_path): """将音频文件转录为文本""" if not asr_pipeline: return "[Error: ASR model not loaded]", "" # 返回错误信息和空路径 if audio_file_path is None: return "", "" # 没有音频输入 print(f"Transcribing audio file: {audio_file_path}") try: # 转录音频 transcription = asr_pipeline(audio_file_path)["text"] print(f"Transcription result: {transcription}") return transcription, audio_file_path # 返回文本和路径（可能用于显示） except Exception as e: print(f"Error during audio transcription: {e}") return f"[Error: Transcription failed: {e}]", audio_file_path # ---- Gradio 应用流程 ---- def process_input(input_text, audio_file, style_choice, quality_choice, neg_prompt, guidance, steps): """处理输入（文本或语音），生成提示词和图像""" final_text_input = "" transcription_source = "" # 用于标记来源 # 优先使用文本框输入 if input_text and input_text.strip(): final_text_input = input_text.strip() transcription_source = " (from text input)" # 如果文本框为空，且有音频文件，则使用语音输入 elif audio_file is not None: transcribed_text, _ = transcribe_audio(audio_file) if transcribed_text and "[Error:" not in transcribed_text: final_text_input = transcribed_text transcription_source = " (from audio input)" elif "[Error:" in transcribed_text: # 如果语音识别出错，直接返回错误信息 return transcribed_text, None # 返回错误提示，不生成图像 else: # 音频为空或识别为空 return "[Error: Please provide input via text or voice]", None else: # 没有有效输入 return "[Error: Please provide input via text or voice]", None print(f"Using input: '{final_text_input}'{transcription_source}") # Step 1: Enhance prompt enhanced_prompt = enhance_prompt(final_text_input, style_modifier=style_choice, quality_boost=quality_choice) print(f"Enhanced prompt: {enhanced_prompt}") # Step 2: Generate image (如果提示词增强成功) generated_image = None if "[Error:" not in enhanced_prompt: try: generated_image = generate_image(enhanced_prompt, neg_prompt, guidance, steps) except gr.Error as e: # 如果 generate_image 抛出 gr.Error，将其信息作为 enhanced_prompt 返回给UI enhanced_prompt = f"{enhanced_prompt}\n\n[Image Generation Error: {e}]" # 不再尝试显示图片 except Exception as e: # 捕获其他意外错误 enhanced_prompt = f"{enhanced_prompt}\n\n[Unexpected Image Generation Error: {e}]" # 返回结果给Gradio界面 return enhanced_prompt, generated_image # ---- Gradio 界面构建 (Step 3: Controls & Step 4: Layout) ---- # 定义可选的风格和质量提升选项 (用于Dropdown/Radio) style_options = ["cinematic", "photorealistic", "anime", "fantasy art", "cyberpunk", "steampunk", "watercolor"] quality_options = ["highly detailed", "sharp focus", "intricate details", "4k", "masterpiece", "best quality"] with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# AI Image Generator: From Idea to Image") gr.Markdown("Enter a short description (or use voice input), and the app will enhance it into a detailed prompt and generate an image using Stable Diffusion.") with gr.Row(): with gr.Column(scale=1): # 输入区域 inp_text = gr.Textbox(label="Enter short description here", placeholder="e.g., A magical treehouse in the sky") # 加分项：语音输入控件 inp_audio = gr.Audio(sources=["microphone"], type="filepath", label="Or record your idea (clears text box if used)", visible=asr_pipeline is not None) # 只有ASR加载成功才显示 # Step 3: 使用不同控件 # 控件1: Dropdown (下拉菜单) inp_style = gr.Dropdown(label="Choose Base Style", choices=style_options, value="cinematic") # 控件2: Radio (单选框) - 也可以用 CheckboxGroup 实现多选 inp_quality = gr.Radio(label="Quality Boost", choices=quality_options, value="highly detailed") # 控件3: Textbox (用于Negative Prompt) inp_neg_prompt = gr.Textbox(label="Negative Prompt (optional)", placeholder="e.g., blurry, low quality, text, watermark") # 控件4: Slider (滑块) inp_guidance = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=7.5, label="Guidance Scale (CFG)") # 控件5: Slider (滑块) inp_steps = gr.Slider(minimum=10, maximum=100, step=1, value=30, label="Inference Steps") # 提交按钮 btn_generate = gr.Button("Generate Image", variant="primary") with gr.Column(scale=1): # 输出区域 out_prompt = gr.Textbox(label="Generated Prompt", interactive=False) # 输出文本框不可编辑 out_image = gr.Image(label="Generated Image", type="pil") # 输出图像 # 设置按钮点击事件 btn_generate.click( fn=process_input, inputs=[inp_text, inp_audio, inp_style, inp_quality, inp_neg_prompt, inp_guidance, inp_steps], outputs=[out_prompt, out_image] ) # (可选) 当用户录音后，可以自动清空文本框，以明确优先使用语音 if asr_pipeline: def clear_text_on_audio(audio_data): if audio_data is not None: return "" # 返回空字符串清空文本框 return gr.update() # 否则不改变文本框内容 (gr.update()是占位符) inp_audio.change(fn=clear_text_on_audio, inputs=inp_audio, outputs=inp_text) # ---- 启动应用 ---- if __name__ == "__main__": # 设置Hugging Face Hub Token (如果需要从私有仓库加载模型) # from huggingface_hub import login # login("YOUR_HF_TOKEN") # 在本地运行时取消注释并替换 # 在Hugging Face Spaces上运行时，端口通常由平台管理 # share=True 会创建一个公共链接 (如果在本地运行需要) demo.launch(share=False)