Update app.py

app.py

@@ -5,9 +5,9 @@ from PIL import Image
 import os
 import gc
 import time
+import spaces
 from typing import Optional, Tuple
 from huggingface_hub import hf_hub_download
-import requests
 
 # Global pipeline variables
 txt2img_pipe = None
@@ -17,7 +17,6 @@ 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"""
@@ -25,130 +24,67 @@ def clear_memory():
         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"""
+    """Load both text2img and img2img pipelines optimized for Spaces"""
     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...")
+    try:
+        print("Loading CyberRealistic Pony models...")
+        
+        # Use Hugging Face Hub download with minimal local storage
+        print(f"Accessing model from {MODEL_REPO}...")
+        
+        # Load Text2Img pipeline
+        if txt2img_pipe is None:
             txt2img_pipe = StableDiffusionXLPipeline.from_single_file(
-                model_path,
+                f"https://huggingface.co/{MODEL_REPO}/resolve/main/{MODEL_FILENAME}",
                 torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                 use_safetensors=True,
                 variant="fp16" if device == "cuda" else None
             )
             
-            # Memory optimizations
+            # Aggressive memory optimizations for Spaces
             txt2img_pipe.enable_attention_slicing()
+            txt2img_pipe.enable_vae_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)
+                txt2img_pipe.enable_model_cpu_offload()
+                txt2img_pipe.enable_sequential_cpu_offload()
             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
+        
+        # Share components for Img2Img to save memory
+        if img2img_pipe is None:
+            img2img_pipe = StableDiffusionXLImg2ImgPipeline(
+                vae=txt2img_pipe.vae,
+                text_encoder=txt2img_pipe.text_encoder,
+                text_encoder_2=txt2img_pipe.text_encoder_2,
+                tokenizer=txt2img_pipe.tokenizer,
+                tokenizer_2=txt2img_pipe.tokenizer_2,
+                unet=txt2img_pipe.unet,
+                scheduler=txt2img_pipe.scheduler,
             )
             
-            # Memory optimizations
+            # Same optimizations
             img2img_pipe.enable_attention_slicing()
+            img2img_pipe.enable_vae_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
+                img2img_pipe.enable_model_cpu_offload()
+                img2img_pipe.enable_sequential_cpu_offload()
+        
+        print("Models loaded successfully!")
+        return True
+        
+    except Exception as e:
+        print(f"Error loading models: {e}")
+        return False
 
 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
         
@@ -157,37 +93,35 @@ def enhance_prompt(prompt: str, add_quality_tags: bool = True) -> str:
 
 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))
+    # More conservative limits for Spaces
+    width = max(512, min(1024, width))
+    height = max(512, min(1024, height))
     
     return width, height
 
+@spaces.GPU(duration=60)  # GPU decorator for Spaces
 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"""
+    """Generate image from text prompt with Spaces GPU support"""
     global txt2img_pipe
     
     if not prompt.strip():
         return None, "Please enter a prompt"
-        
-    # Load models if not already loaded
+    
+    # Lazy load models
     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."
+        if not load_models():
+            return None, "Failed to load models. Please try again."
     
     try:
-        # Clear memory before generation
         clear_memory()
         
-        # Validate and fix dimensions
+        # Validate dimensions
         width, height = validate_dimensions(width, height)
         
-        # Set seed for reproducibility
+        # Set seed
         generator = None
         if seed != -1:
             generator = torch.Generator(device=device).manual_seed(int(seed))
@@ -195,15 +129,15 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
         # Enhance prompt
         enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
         
-        print(f"Generating with prompt: {enhanced_prompt[:100]}...")
+        print(f"Generating: {enhanced_prompt[:100]}...")
         start_time = time.time()
         
-        # Generate image
+        # Generate with lower memory usage
         with torch.no_grad():
             result = txt2img_pipe(
                 prompt=enhanced_prompt,
                 negative_prompt=negative_prompt or "",
-                num_inference_steps=int(num_steps),
+                num_inference_steps=min(int(num_steps), 30),  # Limit steps for Spaces
                 guidance_scale=float(guidance_scale),
                 width=width,
                 height=height,
@@ -211,38 +145,35 @@ def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width,
             )
         
         generation_time = time.time() - start_time
-        status = f"Text2Img: Generated successfully in {generation_time:.1f}s (Size: {width}x{height})"
+        status = f"Generated in {generation_time:.1f}s ({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
+        return None, f"Generation failed: {str(e)}"
     finally:
         clear_memory()
 
+@spaces.GPU(duration=60)  # GPU decorator for Spaces
 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"""
+    """Generate image from input image + text prompt with Spaces GPU support"""
     global img2img_pipe
     
     if input_image is None:
-        return None, "Please upload an input image for Img2Img"
+        return None, "Please upload an input image"
     
     if not prompt.strip():
         return None, "Please enter a prompt"
-        
-    # Load models if not already loaded
+    
+    # Lazy load models
     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."
+        if not load_models():
+            return None, "Failed to load models. Please try again."
     
     try:
-        # Clear memory before generation
         clear_memory()
         
-        # Set seed for reproducibility
+        # Set seed
         generator = None
         if seed != -1:
             generator = torch.Generator(device=device).manual_seed(int(seed))
@@ -252,230 +183,136 @@ def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_s
         
         # 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
+            
+            # Conservative resize for Spaces
+            max_size = 768
             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]}...")
+        print(f"Transforming: {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),
+                num_inference_steps=min(int(num_steps), 30),  # Limit 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})"
+        status = f"Transformed 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
+        return None, f"Transformation failed: {str(e)}"
     finally:
         clear_memory()
 
-# Default negative prompt (improved)
+# Simplified negative prompt for better performance
 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.
+(low quality:1.3), (worst quality:1.3), (bad quality:1.2), blurry, noisy, ugly, deformed, 
+(text, watermark:1.4), (extra limbs:1.3), (bad hands:1.3), (bad anatomy:1.2)
 """
 
-# Create Gradio interface with enhanced styling
+# Gradio interface optimized for Spaces
 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;
-    }
-    """
+    title="CyberRealistic Pony Generator", 
+    theme=gr.themes.Soft()
 ) as demo:
     gr.Markdown("""
-    # 🎨 CyberRealistic Pony Image Generator (Hugging Face Edition)
+    # 🎨 CyberRealistic Pony Image Generator
     
-    Generate high-quality images using the CyberRealistic Pony SDXL model from Hugging Face.
+    Generate high-quality images using the CyberRealistic Pony SDXL model.
     
-    **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).
+    ⚠️ **Note**: First generation may take longer as the model loads. GPU time is limited on Spaces.
     """)
     
     with gr.Tabs():
-        # Text2Image Tab
         with gr.TabItem("🎨 Text to Image"):
             with gr.Row():
-                with gr.Column(scale=1):
-                    # Input controls for Text2Img
+                with gr.Column():
                     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
+                        placeholder="beautiful landscape, mountains, sunset",
+                        lines=2
                     )
                     
-                    with gr.Row():
-                        txt2img_steps = gr.Slider(
-                            minimum=10,
-                            maximum=50,
-                            value=25,
-                            step=1,
-                            label="Inference Steps"
+                    with gr.Accordion("Advanced Settings", open=False):
+                        txt2img_negative = gr.Textbox(
+                            label="Negative Prompt",
+                            value=DEFAULT_NEGATIVE,
+                            lines=2
                         )
                         
-                        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_quality_tags = gr.Checkbox(
+                            label="Add Quality Tags",
+                            value=True
                         )
                         
-                        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.Row():
+                            txt2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
+                            txt2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
+                        
+                        with gr.Row():
+                            txt2img_width = gr.Slider(512, 1024, 768, step=64, label="Width")
+                            txt2img_height = gr.Slider(512, 1024, 768, step=64, label="Height")
+                        
+                        txt2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
                     
-                with gr.Column(scale=2):
-                    # Output for Text2Img
-                    txt2img_output = gr.Image(
-                        label="Generated Image", 
-                        type="pil",
-                        height=600
-                    )
+                    txt2img_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
+                
+                with gr.Column():
+                    txt2img_output = gr.Image(label="Generated Image", height=400)
                     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
-                    )
+                with gr.Column():
+                    img2img_input = gr.Image(label="Input Image", type="pil", height=250)
                     
                     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
+                        placeholder="digital painting style, vibrant colors",
+                        lines=2
                     )
                     
-                    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"
+                    with gr.Accordion("Advanced Settings", open=False):
+                        img2img_negative = gr.Textbox(
+                            label="Negative Prompt",
+                            value=DEFAULT_NEGATIVE,
+                            lines=2
                         )
                         
-                        img2img_guidance = gr.Slider(
-                            minimum=1.0,
-                            maximum=20.0,
-                            value=7.5,
-                            step=0.5,
-                            label="Guidance Scale"
+                        img2img_quality_tags = gr.Checkbox(
+                            label="Add Quality Tags",
+                            value=True
                         )
+                        
+                        with gr.Row():
+                            img2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
+                            img2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
+                        
+                        img2img_strength = gr.Slider(
+                            0.1, 1.0, 0.75, step=0.05, 
+                            label="Strength (Higher = more creative)"
+                        )
+                        
+                        img2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
                     
-                    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_btn = gr.Button("🖼️ Transform", variant="primary", size="lg")
+                
+                with gr.Column():
+                    img2img_output = gr.Image(label="Generated Image", height=400)
                     img2img_status = gr.Textbox(label="Status", interactive=False)
     
     # Event handlers
@@ -488,37 +325,12 @@ with gr.Blocks(
     
     img2img_btn.click(
         fn=generate_img2img,
-        inputs=[img2img_input, img2img_prompt, img2img_negative, txt2img_steps, img2img_guidance, 
+        inputs=[img2img_input, img2img_prompt, img2img_negative, img2img_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()
+print(f"🚀 CyberRealistic Pony Generator initialized on {device}")
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )
+    demo.launch()