Spaces:
Sleeping
Sleeping
| import spaces | |
| import torch | |
| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel | |
| import gradio as gr | |
| from PIL import Image | |
| # 1. 加载 MLSD ControlNet（更适合建筑平面图） | |
| controlnet = ControlNetModel.from_pretrained( | |
| "lllyasviel/sd-controlnet-mlsd", torch_dtype=torch.float16 | |
| ) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| controlnet=controlnet, | |
| torch_dtype=torch.float16, | |
| safety_checker=None, | |
| ).to("cuda") | |
| # 2. 四个视角的 prompt 模板 | |
| VIEW_PROMPTS = { | |
| "前视图（朝管理室）": "standing at the corridor entrance, looking toward the management room, frosted glass window at the end, 2D horror animation style, cold green fluorescent light, terrazzo floor, slight perspective, eye level", | |
| "后视图（朝前门）": "standing at the end of the corridor, looking toward the front entrance, double glass doors, warm street light, red exit sign, 2D horror animation style, slight perspective, eye level", | |
| "左视图（朝楼梯间）": "standing in the middle of the corridor, looking at the left wall, dark gray fire door, green emergency light leaking from the gap, ventilation grille, 2D horror animation style, flat frontal view", | |
| "右视图（朝电梯）": "standing in the middle of the corridor, looking at the right wall, stainless steel elevator doors, red LED floor display, button panel with orange light, 2D horror animation style, flat frontal view", | |
| } | |
| # 3. 用 @spaces.GPU 装饰器包裹生成函数 | |
| def generate_image(control_image, view_choice, extra_prompt): | |
| prompt = VIEW_PROMPTS[view_choice] | |
| if extra_prompt: | |
| prompt += ", " + extra_prompt | |
| result = pipe( | |
| prompt=prompt, | |
| image=control_image, | |
| num_inference_steps=30, | |
| controlnet_conditioning_scale=0.8, | |
| ).images[0] | |
| return result | |
| # 4. Gradio 界面：一次选一个视角 | |
| iface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.Image(type="pil", label="上传平面图（ControlNet 输入）"), | |
| gr.Dropdown( | |
| choices=list(VIEW_PROMPTS.keys()), | |
| value="前视图（朝管理室）", | |
| label="选择视角", | |
| ), | |
| gr.Textbox(label="附加风格描述（可选）", placeholder="例如：flickering light, dark handprint on window"), | |
| ], | |
| outputs=gr.Image(label="生成结果"), | |
| title="公寓走廊 · 四视角生成器", | |
| description="上传平面图，选择视角，生成对应的场景图。每次生成一个视角，跑 4 次得到完整四视图。", | |
| ) | |
| iface.launch() |