Create app.py

app.py

@@ -0,0 +1,193 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import DDIMScheduler
+from PIL import Image
+import cv2
+import numpy as np
+from controlnet_aux import CannyDetector
+import os
+
+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
+        # )
+
+        # Load Realistic LoRA model
+        self.pipe = StableDiffusionControlNetPipeline.from_single_file(
+            "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):
+        """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=20, guidance_scale=7.0, 
+                               controlnet_conditioning_scale=1.0, 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!")
+
+def generate_face(sketch_image, custom_prompt, seed):
+    """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
+        )
+        
+        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. Set a seed for reproducible results (optional)
+        4. 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
+            )
+            
+            seed_input = gr.Number(
+                label="Seed (optional)",
+                placeholder="Enter a number for reproducible results",
+                precision=0
+            )
+            
+            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:
+        - 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], outputs=[realistic_output, control_output])
+
+# Launch the app
+if __name__ == "__main__":
+    app.launch(share=True)