Spaces:
Sleeping
Sleeping
| 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" | |
| # Determine device | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {device}") | |
| # Load the entire pipeline directly from the Hugging Face Hub | |
| # The LoRA weights are already integrated into this saved pipeline | |
| try: | |
| # Attempt to load custom model from Hub | |
| controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16 if device == "cuda" else torch.float32) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| HF_REPO_ID, | |
| controlnet=controlnet, | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32, | |
| low_cpu_mem_usage=False # Ensure weights are fully loaded and not in meta device | |
| ) | |
| 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 if device == "cuda" else torch.float32) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| controlnet=controlnet, | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32, | |
| low_cpu_mem_usage=False # Ensure weights are fully loaded and not in meta device | |
| ) | |
| print("Using base Stable Diffusion v1.5 with ControlNet.") | |
| pipe.to(device) | |
| 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 | |