Create app.py

app.py

@@ -0,0 +1,281 @@
+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)