File size: 2,502 Bytes
11be1a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from controlnet_aux import CannyDetector
from PIL import Image
import gradio as gr
import os
# Define the Hugging Face repository ID for your model
HF_REPO_ID = "saad206121/ai_interior_design_model"
# Load the entire pipeline directly from the Hugging Face Hub
# The LoRA weights are already integrated into this saved pipeline
try:
pipe = StableDiffusionControlNetPipeline.from_pretrained(HF_REPO_ID, torch_dtype=torch.float16)
print(f"✅ Model loaded successfully from Hugging Face Hub: {HF_REPO_ID}")
except Exception as e:
print(f"⚠️ Could not load model from Hugging Face Hub ({HF_REPO_ID}): {e}. Please check the repo ID and permissions.")
# Fallback to loading base model and controlnet if custom model fails
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)
print("Using base Stable Diffusion v1.5 with ControlNet.")
pipe.to("cuda")
canny = CannyDetector()
def design_room(image, user_prompt, negative_prompt, num_steps, guidance_scale):
if image is None:
return None
# Preprocess: Get edges
image = Image.fromarray(image)
canny_image = canny(image, detect_resolution=512, image_resolution=512)
# Combine user prompt with trigger word
full_prompt = f"sks bedroom interior, {user_prompt}"
result = pipe(
full_prompt,
image=canny_image,
negative_prompt=negative_prompt,
num_inference_steps=num_steps,
guidance_scale=guidance_scale
).images[0]
return result
# --- GRADIO UI ---
iface = gr.Interface(
fn=design_room,
inputs=[
gr.Image(label="Upload Room Photo"),
gr.Textbox(label="Prompt (e.g. 'modern luxury, blue furniture')"),
gr.Textbox(label="Negative Prompt", value="low quality, blurry, distorted, messy"),
gr.Slider(10, 50, value=50, label="Steps"),
gr.Slider(1, 15, value=15, label="Guidance Scale")
],
outputs=gr.Image(label="Designed Room"),
title="AI Interior Designer (Custom Trained)",
description="Upload a photo of a room. The AI will redesign it using your custom trained style while keeping the furniture in the same place."
)
iface.launch(share=False) # share=False for Hugging Face Spaces
|