app.py

import gradio as gr
import spaces

import os
import cv2
import numpy as np
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import DDIMScheduler
from controlnet_aux import CannyDetector
from PIL import Image
import base64
from io import BytesIO


class SketchToRealisticFace:
    def __init__(self):
        """Initialize the sketch-to-realistic face pipeline"""
        
        # Load ControlNet model for Canny edge detection
        self.controlnet = ControlNetModel.from_pretrained(
            "lllyasviel/sd-controlnet-canny",
            torch_dtype=torch.float16
        )
        
        # # Load base SD 1.5 model
        # self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
        #     "runwayml/stable-diffusion-v1-5",
        #     controlnet=self.controlnet,
        #     torch_dtype=torch.float16,
        #     safety_checker=None,
        #     requires_safety_checker=False
        # )

        self.pipe = StableDiffusionControlNetPipeline.from_single_file(
            "https://huggingface.co/NikhilJoson/Realistic_Vision_V6_B1/blob/main/realisticVisionV60B1_v51HyperVAE.safetensors",
            controlnet=self.controlnet,
            torch_dtype=torch.float16,
            safety_checker=None,
            requires_safety_checker=False
        )
        
        # Use DDIM scheduler for better quality
        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
        
        # Move to GPU if available
        if torch.cuda.is_available():
            self.pipe = self.pipe.to("cuda")
        
        # Initialize Canny detector
        self.canny_detector = CannyDetector()
        
        # Enable memory efficient attention
        #self.pipe.enable_memory_efficient_attention()
        
        # Set default parameters
        self.default_prompt = "RAW photo, portrait, 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3"
        self.default_negative_prompt = "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime), text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck"
        
    def preprocess_sketch(self, sketch_image, low_threshold=100, high_threshold=200):
        """Preprocess sketch image to create Canny edge map"""
        image_array = np.array(sketch_image)
        
        if len(image_array.shape) == 3:
            image_array = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
        
        canny = cv2.Canny(image_array, low_threshold, high_threshold)
        control_image = Image.fromarray(canny).convert("RGB")
        
        return control_image
    
    def generate(self, sketch_image, custom_prompt=None, width=512, height=512, seed=None, 
                 num_inference_steps=20, guidance_scale=7.0, controlnet_conditioning_scale=1.0):
        """Generate realistic face from sketch"""
        
        if sketch_image is None:
            return None, None
        
        # Use custom prompt if provided, otherwise use default
        prompt = custom_prompt if custom_prompt and custom_prompt.strip() else self.default_prompt
        
        # Resize sketch to target dimensions
        sketch_image = sketch_image.resize((width, height))
        
        # Preprocess sketch to create control image
        control_image = self.preprocess_sketch(sketch_image)
        
        # Set seed for reproducibility
        generator = torch.Generator(device=self.pipe.device).manual_seed(seed) if seed else None
        
        # Generate image
        with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
            result = self.pipe(prompt=prompt, image=control_image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, 
                               controlnet_conditioning_scale=controlnet_conditioning_scale, generator=generator, width=width, height=height,
                               negative_prompt=self.default_negative_prompt)
        
        return result.images[0], control_image

# Initialize the generator globally
print("Loading model... This may take a few minutes.")
generator = SketchToRealisticFace()
print("Model loaded successfully!")

@spaces.GPU
def generate_face(sketch_image, custom_prompt, seed, num_inference_steps, guidance_scale, controlnet_conditioning_scale):
    """Wrapper function for Gradio interface"""
    try:
        # Convert seed to int if provided
        seed_int = int(seed) if seed else None
        
        # Generate the realistic face
        realistic_face, control_image = generator.generate(sketch_image=sketch_image, custom_prompt=custom_prompt, seed=seed_int,
                                                           num_inference_steps=num_inference_steps, guidance_scale=guidance_scale,
                                                           controlnet_conditioning_scale=controlnet_conditioning_scale)
        
        return realistic_face, control_image
        
    except Exception as e:
        print(f"Error: {str(e)}")
        return None, None

# Create Gradio interface
with gr.Blocks(title="Sketch to Realistic Face Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🎨 Sketch to Realistic Face Generator
        
        Transform your sketches into realistic faces using Stable Diffusion with ControlNet!
        
        **Instructions:**
        1. Upload a sketch or drawing of a face
        2. Optionally customize the prompt
        3. Adjust generation parameters (steps, guidance scale, etc.)
        4. Set a seed for reproducible results (optional)
        5. Click "Generate Realistic Face"
        """
    )
    
    with gr.Row():
        with gr.Column():
            # Input components
            sketch_input = gr.Image(label="Upload Sketch", type="pil", height=400)
            
            custom_prompt = gr.Textbox(label="Custom Prompt (optional)",
                                       placeholder="Leave empty to use default prompt, or customize: 'portrait of a young person, professional headshot, studio lighting...'",
                                       lines=3)
            
            with gr.Row():
                seed_input = gr.Number(label="Seed (optional)", placeholder="Enter a number for reproducible results", precision=0)
            
            # Generation parameters
            gr.Markdown("### 🎛️ Generation Parameters")
            
            num_inference_steps = gr.Slider(minimum=10, maximum=50, value=20, step=1, label="Inference Steps",
                                            info="More steps = higher quality but slower generation")
            
            guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=7.0, step=0.5, label="Guidance Scale",
                                       info="How closely to follow the prompt (higher = more adherence to prompt)")
            
            controlnet_conditioning_scale = gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="ControlNet Conditioning Scale",
                                                      info="How strongly to follow the sketch structure")
            
            generate_btn = gr.Button("🚀 Generate Realistic Face", variant="primary", size="lg")
        
        with gr.Column():
            # Output components
            with gr.Row():
                realistic_output = gr.Image(label="Generated Realistic Face", height=400)
                control_output = gr.Image(label="Control Image (Canny Edges)", height=400)
    
    # Add examples
    gr.Markdown("## 📝 Default Prompt")
    gr.Markdown(f"```{generator.default_prompt}```")
    
    gr.Markdown(
        """
        ## 💡 Tips:
        - **Inference Steps**: 20-30 steps usually provide good quality. More steps improve quality but increase generation time.
        - **Guidance Scale**: 7.0-12.0 works well. Higher values make the AI follow your prompt more strictly.
        - **ControlNet Scale**: 1.0 is usually perfect. Lower values give more creative freedom, higher values stick closer to sketch structure.
        - Upload clear sketches with well-defined facial features
        - The model works best with front-facing portraits
        - Use the same seed number to get consistent results
        - Customize the prompt to specify style, lighting, or other details
        - The control image shows how your sketch is interpreted as edges
        """
    )
    
    # Connect the function to the interface
    generate_btn.click(
        fn=generate_face, 
        inputs=[sketch_input, custom_prompt, seed_input, num_inference_steps, guidance_scale, controlnet_conditioning_scale], 
        outputs=[realistic_output, control_output]
    )

# Launch the app
if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, share=True)