Create app.py

app.py

@@ -0,0 +1,524 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+from PIL import Image
+import os
+import gc
+import time
+from typing import Optional, Tuple
+from huggingface_hub import hf_hub_download
+import requests
+
+# Global pipeline variables
+txt2img_pipe = None
+img2img_pipe = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Hugging Face model configuration
+MODEL_REPO = "ajsbsd/CyberRealistic-Pony"
+MODEL_FILENAME = "cyberrealisticPony_v110.safetensors"
+LOCAL_MODEL_PATH = "./models/cyberrealisticPony_v110.safetensors"
+
+def clear_memory():
+    """Clear GPU memory"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+def download_model():
+    """Download model from Hugging Face if not already cached"""
+    try:
+        # Create models directory if it doesn't exist
+        os.makedirs("./models", exist_ok=True)
+        
+        # Check if model already exists locally
+        if os.path.exists(LOCAL_MODEL_PATH):
+            print(f"Model already exists at {LOCAL_MODEL_PATH}")
+            return LOCAL_MODEL_PATH
+        
+        print(f"Downloading model from {MODEL_REPO}/{MODEL_FILENAME}...")
+        print("This may take a while on first run...")
+        
+        # Download the model file
+        model_path = hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILENAME,
+            local_dir="./models",
+            local_dir_use_symlinks=False,
+            resume_download=True
+        )
+        
+        print(f"Model downloaded successfully to {model_path}")
+        return model_path
+        
+    except Exception as e:
+        print(f"Error downloading model: {e}")
+        print("Attempting to use cached version or fallback...")
+        
+        # Try to use Hugging Face cache directly
+        try:
+            cached_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename=MODEL_FILENAME,
+                resume_download=True
+            )
+            print(f"Using cached model at {cached_path}")
+            return cached_path
+        except Exception as cache_error:
+            print(f"Cache fallback failed: {cache_error}")
+            return None
+
+def load_models():
+    """Load both text2img and img2img pipelines with Hugging Face integration"""
+    global txt2img_pipe, img2img_pipe
+    
+    # Download model if needed
+    model_path = download_model()
+    
+    if model_path is None:
+        print("Failed to download or locate model file")
+        return None, None
+    
+    if not os.path.exists(model_path):
+        print(f"Model file not found after download: {model_path}")
+        return None, None
+    
+    if txt2img_pipe is None:
+        try:
+            print("Loading CyberRealistic Pony Text2Img model...")
+            txt2img_pipe = StableDiffusionXLPipeline.from_single_file(
+                model_path,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                use_safetensors=True,
+                variant="fp16" if device == "cuda" else None
+            )
+            
+            # Memory optimizations
+            txt2img_pipe.enable_attention_slicing()
+            
+            if device == "cuda":
+                try:
+                    txt2img_pipe.enable_model_cpu_offload()
+                    print("Text2Img CPU offload enabled")
+                except Exception as e:
+                    print(f"Text2Img CPU offload failed: {e}")
+                    txt2img_pipe = txt2img_pipe.to(device)
+            else:
+                txt2img_pipe = txt2img_pipe.to(device)
+                    
+            print("Text2Img model loaded successfully!")
+            
+        except Exception as e:
+            print(f"Error loading Text2Img model: {e}")
+            return None, None
+    
+    if img2img_pipe is None:
+        try:
+            print("Loading CyberRealistic Pony Img2Img model...")
+            img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
+                model_path,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                use_safetensors=True,
+                variant="fp16" if device == "cuda" else None
+            )
+            
+            # Memory optimizations
+            img2img_pipe.enable_attention_slicing()
+            
+            if device == "cuda":
+                try:
+                    img2img_pipe.enable_model_cpu_offload()
+                    print("Img2Img CPU offload enabled")
+                except Exception as e:
+                    print(f"Img2Img CPU offload failed: {e}")
+                    img2img_pipe = img2img_pipe.to(device)
+            else:
+                img2img_pipe = img2img_pipe.to(device)
+                    
+            print("Img2Img model loaded successfully!")
+            
+        except Exception as e:
+            print(f"Error loading Img2Img model: {e}")
+            return txt2img_pipe, None
+    
+    return txt2img_pipe, img2img_pipe
+
+def enhance_prompt(prompt: str, add_quality_tags: bool = True) -> str:
+    """Enhance prompt with Pony-style tags"""
+    if not prompt.strip():
+        return prompt
+        
+    # Don't add tags if already present
+    if prompt.startswith("score_") or not add_quality_tags:
+        return prompt
+        
+    quality_tags = "score_9, score_8_up, score_7_up, masterpiece, best quality, highly detailed"
+    return f"{quality_tags}, {prompt}"
+
+def validate_dimensions(width: int, height: int) -> Tuple[int, int]:
+    """Ensure dimensions are valid for SDXL"""
+    # SDXL works best with dimensions divisible by 64
+    width = ((width + 63) // 64) * 64
+    height = ((height + 63) // 64) * 64
+    
+    # Ensure reasonable limits
+    width = max(512, min(1536, width))
+    height = max(512, min(1536, height))
+    
+    return width, height
+
+def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed, add_quality_tags):
+    """Generate image from text prompt with enhanced error handling"""
+    global txt2img_pipe
+    
+    if not prompt.strip():
+        return None, "Please enter a prompt"
+        
+    # Load models if not already loaded
+    if txt2img_pipe is None:
+        txt2img_pipe, _ = load_models()
+        if txt2img_pipe is None:
+            return None, "Failed to load Text2Img model. Please check your internet connection and try again."
+    
+    try:
+        # Clear memory before generation
+        clear_memory()
+        
+        # Validate and fix dimensions
+        width, height = validate_dimensions(width, height)
+        
+        # Set seed for reproducibility
+        generator = None
+        if seed != -1:
+            generator = torch.Generator(device=device).manual_seed(int(seed))
+        
+        # Enhance prompt
+        enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
+        
+        print(f"Generating with prompt: {enhanced_prompt[:100]}...")
+        start_time = time.time()
+        
+        # Generate image
+        with torch.no_grad():
+            result = txt2img_pipe(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt or "",
+                num_inference_steps=int(num_steps),
+                guidance_scale=float(guidance_scale),
+                width=width,
+                height=height,
+                generator=generator
+            )
+        
+        generation_time = time.time() - start_time
+        status = f"Text2Img: Generated successfully in {generation_time:.1f}s (Size: {width}x{height})"
+        
+        return result.images[0], status
+        
+    except Exception as e:
+        error_msg = f"Text2Img generation failed: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+    finally:
+        clear_memory()
+
+def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_scale, strength, seed, add_quality_tags):
+    """Generate image from input image + text prompt with enhanced error handling"""
+    global img2img_pipe
+    
+    if input_image is None:
+        return None, "Please upload an input image for Img2Img"
+    
+    if not prompt.strip():
+        return None, "Please enter a prompt"
+        
+    # Load models if not already loaded
+    if img2img_pipe is None:
+        _, img2img_pipe = load_models()
+        if img2img_pipe is None:
+            return None, "Failed to load Img2Img model. Please check your internet connection and try again."
+    
+    try:
+        # Clear memory before generation
+        clear_memory()
+        
+        # Set seed for reproducibility
+        generator = None
+        if seed != -1:
+            generator = torch.Generator(device=device).manual_seed(int(seed))
+        
+        # Enhance prompt
+        enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
+        
+        # Process input image
+        if isinstance(input_image, Image.Image):
+            # Ensure RGB format
+            if input_image.mode != 'RGB':
+                input_image = input_image.convert('RGB')
+                
+            # Resize to reasonable dimensions while maintaining aspect ratio
+            original_size = input_image.size
+            max_size = 1024
+            input_image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+            
+            # Ensure dimensions are divisible by 64
+            w, h = input_image.size
+            w, h = validate_dimensions(w, h)
+            input_image = input_image.resize((w, h), Image.Resampling.LANCZOS)
+        
+        print(f"Generating with prompt: {enhanced_prompt[:100]}...")
+        start_time = time.time()
+        
+        # Generate image
+        with torch.no_grad():
+            result = img2img_pipe(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt or "",
+                image=input_image,
+                num_inference_steps=int(num_steps),
+                guidance_scale=float(guidance_scale),
+                strength=float(strength),
+                generator=generator
+            )
+        
+        generation_time = time.time() - start_time
+        status = f"Img2Img: Generated successfully in {generation_time:.1f}s (Strength: {strength})"
+        
+        return result.images[0], status
+        
+    except Exception as e:
+        error_msg = f"Img2Img generation failed: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+    finally:
+        clear_memory()
+
+# Default negative prompt (improved)
+DEFAULT_NEGATIVE = """
+(low quality:1.4), (worst quality:1.4), (bad quality:1.3), (normal quality:1.2), lowres, jpeg artifacts, blurry, noisy, ugly, deformed, disfigured, malformed, poorly drawn, bad art, amateur, render, 3D, cgi, 
+(text, signature, watermark, username, copyright:1.5), 
+(extra limbs:1.5), (missing limbs:1.5), (extra fingers:1.5), (missing fingers:1.5), (mutated hands:1.5), (bad hands:1.4), (poorly drawn hands:1.3), (ugly hands:1.2), 
+(bad anatomy:1.4), (deformed body:1.3), (unnatural body:1.2), (cross-eyed:1.3), (skewed eyes:1.3), (imperfect eyes:1.2), (ugly eyes:1.2), (asymmetrical face:1.2), (unnatural face:1.2), 
+(blush:1.1), (shadow on skin:1.1), (shaded skin:1.1), (dark skin:1.1), 
+abstract, simplified, unrealistic, impressionistic, cartoon, anime, drawing, sketch, illustration, painting, censored, grayscale, monochrome, out of frame, cropped, distorted.
+"""
+
+# Create Gradio interface with enhanced styling
+with gr.Blocks(
+    title="CyberRealistic Pony Image Generator", 
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+    }
+    .tab-nav button {
+        font-size: 16px;
+        font-weight: bold;
+    }
+    """
+) as demo:
+    gr.Markdown("""
+    # 🎨 CyberRealistic Pony Image Generator (Hugging Face Edition)
+    
+    Generate high-quality images using the CyberRealistic Pony SDXL model from Hugging Face.
+    
+    **Features:**
+    - 🎨 Text-to-Image generation
+    - 🖼️ Image-to-Image transformation
+    - 🎯 Automatic quality tag enhancement
+    - ⚡ Memory optimized for better performance
+    - 🤗 Auto-downloads model from Hugging Face
+    
+    **Note:** On first run, the model will be downloaded from Hugging Face (this may take a few minutes).
+    """)
+    
+    with gr.Tabs():
+        # Text2Image Tab
+        with gr.TabItem("🎨 Text to Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    # Input controls for Text2Img
+                    txt2img_prompt = gr.Textbox(
+                        label="Prompt",
+                        placeholder="Enter your image description...",
+                        value="beautiful landscape with mountains and lake at sunset",
+                        lines=3
+                    )
+                    
+                    txt2img_negative = gr.Textbox(
+                        label="Negative Prompt",
+                        value=DEFAULT_NEGATIVE,
+                        lines=3
+                    )
+                    
+                    txt2img_quality_tags = gr.Checkbox(
+                        label="Add Quality Tags",
+                        value=True
+                    )
+                    
+                    with gr.Row():
+                        txt2img_steps = gr.Slider(
+                            minimum=10,
+                            maximum=50,
+                            value=25,
+                            step=1,
+                            label="Inference Steps"
+                        )
+                        
+                        txt2img_guidance = gr.Slider(
+                            minimum=1.0,
+                            maximum=20.0,
+                            value=7.5,
+                            step=0.5,
+                            label="Guidance Scale"
+                        )
+                    
+                    with gr.Row():
+                        txt2img_width = gr.Slider(
+                            minimum=512,
+                            maximum=1536,
+                            value=1024,
+                            step=64,
+                            label="Width"
+                        )
+                        
+                        txt2img_height = gr.Slider(
+                            minimum=512,
+                            maximum=1536,
+                            value=1024,
+                            step=64,
+                            label="Height"
+                        )
+                    
+                    txt2img_seed = gr.Number(
+                        label="Seed (-1 for random)",
+                        value=-1,
+                        precision=0
+                    )
+                    
+                    txt2img_btn = gr.Button("🎨 Generate Image", variant="primary")
+                    
+                with gr.Column(scale=2):
+                    # Output for Text2Img
+                    txt2img_output = gr.Image(
+                        label="Generated Image", 
+                        type="pil",
+                        height=600
+                    )
+                    txt2img_status = gr.Textbox(label="Status", interactive=False)
+        
+        # Image2Image Tab
+        with gr.TabItem("🖼️ Image to Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    # Input controls for Img2Img
+                    img2img_input = gr.Image(
+                        label="Input Image",
+                        type="pil",
+                        height=300
+                    )
+                    
+                    img2img_prompt = gr.Textbox(
+                        label="Prompt",
+                        placeholder="Describe how to modify the image...",
+                        value="in the style of a digital painting, vibrant colors",
+                        lines=3
+                    )
+                    
+                    img2img_negative = gr.Textbox(
+                        label="Negative Prompt",
+                        value=DEFAULT_NEGATIVE,
+                        lines=3
+                    )
+                    
+                    img2img_quality_tags = gr.Checkbox(
+                        label="Add Quality Tags",
+                        value=True
+                    )
+                    
+                    with gr.Row():
+                        img2img_steps = gr.Slider(
+                            minimum=10,
+                            maximum=50,
+                            value=25,
+                            step=1,
+                            label="Inference Steps"
+                        )
+                        
+                        img2img_guidance = gr.Slider(
+                            minimum=1.0,
+                            maximum=20.0,
+                            value=7.5,
+                            step=0.5,
+                            label="Guidance Scale"
+                        )
+                    
+                    img2img_strength = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.75,
+                        step=0.05,
+                        label="Denoising Strength (Lower = more like input, Higher = more creative)"
+                    )
+                    
+                    img2img_seed = gr.Number(
+                        label="Seed (-1 for random)",
+                        value=-1,
+                        precision=0
+                    )
+                    
+                    img2img_btn = gr.Button("🖼️ Transform Image", variant="primary")
+                    
+                with gr.Column(scale=2):
+                    # Output for Img2Img
+                    img2img_output = gr.Image(
+                        label="Generated Image", 
+                        type="pil",
+                        height=600
+                    )
+                    img2img_status = gr.Textbox(label="Status", interactive=False)
+    
+    # Event handlers
+    txt2img_btn.click(
+        fn=generate_txt2img,
+        inputs=[txt2img_prompt, txt2img_negative, txt2img_steps, txt2img_guidance, 
+                txt2img_width, txt2img_height, txt2img_seed, txt2img_quality_tags],
+        outputs=[txt2img_output, txt2img_status]
+    )
+    
+    img2img_btn.click(
+        fn=generate_img2img,
+        inputs=[img2img_input, img2img_prompt, img2img_negative, txt2img_steps, img2img_guidance, 
+                img2img_strength, img2img_seed, img2img_quality_tags],
+        outputs=[img2img_output, img2img_status]
+    )
+
+# Load models on startup
+print("Initializing CyberRealistic Pony Generator (Hugging Face Edition)...")
+print(f"Device: {device}")
+print(f"Model Repository: {MODEL_REPO}")
+print(f"Model File: {MODEL_FILENAME}")
+
+# Pre-load models in a separate thread to avoid blocking startup
+import threading
+
+def preload_models():
+    """Pre-load models in background"""
+    try:
+        print("Starting background model loading...")
+        load_models()
+        print("Background model loading completed!")
+    except Exception as e:
+        print(f"Background model loading failed: {e}")
+
+# Start background loading
+loading_thread = threading.Thread(target=preload_models, daemon=True)
+loading_thread.start()
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )