sango1/11

1import gradio as gr2import torch3from diffusers import StableDiffusionPipeline4from PIL import Image5import traceback6from typing import Optional7 8# Stable Diffusion模型相关设置9model_id: str = "runwayml/stable-diffusion-v1-5"10device: str = "cpu"  # force CPU usage for compatibility11 12image_generator_pipe: Optional[StableDiffusionPipeline] = None13 14try:15    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)16    image_generator_pipe = pipe.to(device)17except Exception as e:18    print(f"Failed to load Stable Diffusion model: {e}")19 20# 提示词优化函数（简单版）21def optimize_prompt_simple(short_description: str) -> str:22    optimized_prompt = f"Generate a high-quality, detailed image based on the following description: {short_description}"23    return optimized_prompt24 25# 图像生成函数26def generate_image_sd(short_description: str,27                      negative_prompt: str,28                      guidance_scale: float,29                      num_inference_steps: int) -> Image.Image:30    optimized_prompt = optimize_prompt_simple(short_description)31 32    try:33        with torch.no_grad():34            if image_generator_pipe is None:35                raise RuntimeError("Stable Diffusion pipeline is not available.")36 37            output = image_generator_pipe(38                prompt=optimized_prompt,39                negative_prompt=negative_prompt,40                guidance_scale=guidance_scale,41                num_inference_steps=num_inference_steps42            )43            image = output.images[0] if output.images else None44 45        if not image:46            raise RuntimeError("No image was returned from the generation pipeline.")47 48        return image49    except Exception as e:50        traceback.print_exc()51        raise gr.Error(f"Image generation failed: {str(e)}")52    # Gradio界面53with gr.Blocks(theme=gr.themes.Soft()) as demo:54    with gr.Row():55        with gr.Column(scale=1):56            short_description = gr.Textbox(label="Short Description", placeholder="A magical treehouse in the sky")57            optimized_prompt_display = gr.Textbox(label="Optimized Prompt", interactive=False)58            neg_prompt = gr.Textbox(label="Negative Prompt", placeholder="blurry, distorted, watermark")59            guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="Guidance Scale")60            steps = gr.Slider(10, 50, value=25, step=1, label="Inference Steps")61            generate_btn = gr.Button("Generate Image")62 63        with gr.Column(scale=1):64            output_image = gr.Image(label="Generated Image", type="pil")65 66    # 当用户输入简短描述时，自动优化提示词并显示67    short_description.input(68        fn=lambda x: optimize_prompt_simple(x),69        inputs=short_description,70        outputs=optimized_prompt_display71    )72 73    generate_btn.click(74        fn=generate_image_sd,75        inputs=[short_description, neg_prompt, guidance, steps],76        outputs=output_image77    )78 79if __name__ == "__main__":80    if not image_generator_pipe:81        print("WARNING: Stable Diffusion pipeline is not available. UI will launch, but generation will fail.")82 83    demo.launch(server_name="0.0.0.0", server_port=7860)