Update app.py

app.py

@@ -8,37 +8,27 @@ from diffusers import StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
 import torch
 from compel import Compel, ReturnedEmbeddingsType
-#import os
-#from gradio_client import Client
-#client = Client("dhead/ntr-mix-illustrious-xl-noob-xl-xiii-sdxl", hf_token=os.getenv("HUGGING_FACE_TOKEN"))
+import gc
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Check if CUDA is available
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA device: {torch.cuda.get_device_name(0)}")
 
-# Make sure to use torch.float16 consistently throughout the pipeline
+# Initialize the pipeline with CPU first to avoid CUDA errors during model loading
 pipe = StableDiffusionXLPipeline.from_pretrained(
     "dhead/wai-nsfw-illustrious-sdxl-v140-sdxl",
     torch_dtype=torch.float16,
-    variant="fp16",  # Explicitly use fp16 variant
-    use_safetensors=True  # Use safetensors if available
+    variant="fp16",
+    use_safetensors=True,
+    device_map=None,  # Don't auto-assign devices
+    low_cpu_mem_usage=True  # Optimize CPU memory usage
 )
 
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-pipe.to(device)
 
-# Force all components to use the same dtype
-pipe.text_encoder.to(torch.float16)
-pipe.text_encoder_2.to(torch.float16)
-pipe.vae.to(torch.float16)
-pipe.unet.to(torch.float16)
-
-# 追加: Initialize Compel for long prompt processing
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True],
-    truncate_long_prompts=False
-)
+# Initialize Compel for long prompt processing (will be moved to GPU when needed)
+compel = None  # Will be initialized in the GPU function
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1216
@@ -46,28 +36,50 @@ MAX_IMAGE_SIZE = 1216
 # Default prompt
 DEFAULT_PROMPT = "Detailed illustration, realistic style, portrait of a beautiful Japanese woman, wearing an elegant traditional Japanese uniform, neatly tailored with intricate patterns and subtle textures, serene expression, soft natural lighting, standing gracefully in a traditional Japanese garden with cherry blossom petals gently falling in the background, cinematic quality, ultra-detailed, high-resolution, warm tones"
 
-# 追加: Simple long prompt processing function
+def initialize_compel():
+    """Initialize Compel with the current pipeline's components"""
+    global compel
+    if compel is None:
+        compel = Compel(
+            tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+            text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+            requires_pooled=[False, True],
+            truncate_long_prompts=False
+        )
+    return compel
+
 def process_long_prompt(prompt, negative_prompt=""):
     """Simple long prompt processing using Compel"""
     try:
-        conditioning, pooled = compel([prompt, negative_prompt])
+        comp = initialize_compel()
+        conditioning, pooled = comp([prompt, negative_prompt])
         return conditioning, pooled
     except Exception as e:
         print(f"Long prompt processing failed: {e}, falling back to standard processing")
         return None, None
-    
-@spaces.GPU
+
+@spaces.GPU(duration=60)  # Add duration parameter for longer generation times
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    # 変更: Remove the 60-word limit warning and add long prompt check
+    # Move pipeline to GPU inside the GPU-decorated function
+    global pipe
+    pipe = pipe.to("cuda")
+    
+    # Ensure all components are on GPU with correct dtype
+    pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float16)
+    pipe.text_encoder_2 = pipe.text_encoder_2.to(dtype=torch.float16)
+    pipe.vae = pipe.vae.to(dtype=torch.float16)
+    pipe.unet = pipe.unet.to(dtype=torch.float16)
+    
     use_long_prompt = len(prompt.split()) > 60 or len(prompt) > 300
         
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     
     try:
-        # 追加: Try long prompt processing first if prompt is long
+        # Try long prompt processing first if prompt is long
         if use_long_prompt:
             print("Using long prompt processing...")
             conditioning, pooled = process_long_prompt(prompt, negative_prompt)
@@ -101,184 +113,9 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     except RuntimeError as e:
         print(f"Error during generation: {e}")
         # Return a blank image with error message
-        error_img = Image.new('RGB', (width, height), color=(0, 0, 0))
+        error_img = Image.new('RGB', (width, height), color=(50, 50, 50))
         return error_img
-
-
-css = """
-/* Main container styling */
-#col-container {
-    margin: 0 auto;
-    max-width: 1024px;
-}
-
-/* Gradient background for the entire app */
-.gradio-container {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #f5576c 75%, #ffc947 100%);
-    min-height: 100vh;
-}
-
-/* Main block styling with semi-transparent background */
-.contain {
-    background: rgba(255, 255, 255, 0.95);
-    border-radius: 20px;
-    padding: 20px;
-    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
-    backdrop-filter: blur(4px);
-    border: 1px solid rgba(255, 255, 255, 0.18);
-}
-
-/* Input field styling */
-.gr-text-input {
-    background: rgba(255, 255, 255, 0.9) !important;
-    border: 2px solid rgba(102, 126, 234, 0.3) !important;
-    border-radius: 10px !important;
-}
-
-/* Button styling */
-.gr-button {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
-    border: none !important;
-    color: white !important;
-    font-weight: bold !important;
-    transition: all 0.3s ease !important;
-}
-
-.gr-button:hover {
-    transform: translateY(-2px);
-    box-shadow: 0 5px 15px rgba(102, 126, 234, 0.4);
-}
-
-/* Accordion styling */
-.gr-accordion {
-    background: rgba(255, 255, 255, 0.8) !important;
-    border-radius: 10px !important;
-    margin-top: 10px !important;
-}
-
-/* Result image container */
-.gr-image {
-    border-radius: 15px !important;
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1) !important;
-}
-
-/* Slider styling */
-.gr-slider {
-    background: rgba(255, 255, 255, 0.8) !important;
-}
-
-/* Additional styling for headers */
-h1, h2, h3 {
-    color: #333 !important;
-    text-align: center;
-}
-
-/* Markdown text styling */
-.gr-markdown {
-    text-align: center;
-    margin-bottom: 20px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            """
-            # 🎨 Stable Diffusion XL Image Generator
-            ### Create stunning AI-generated images with advanced controls
-            """
-        )
-
-        # Badge section
-        gr.HTML(
-            """
-            <div style="display: flex; justify-content: center; align-items: center; gap: 20px; margin: 20px 0;">
-                <a href="https://huggingface.co/spaces/Heartsync/Wan-2.2-ADULT" target="_blank">
-                    <img src="https://img.shields.io/static/v1?label=T2I%20%26%20TI2V&message=Wan-2.2-ADULT&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
-                </a>
-                <a href="https://huggingface.co/spaces/Heartsync/PornHUB" target="_blank">
-                    <img src="https://img.shields.io/static/v1?label=T2I%20&message=PornHUB&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
-                </a>
-                <a href="https://huggingface.co/spaces/Heartsync/Hentai-Adult" target="_blank">
-                    <img src="https://img.shields.io/static/v1?label=T2I%20&message=Hentai-Adult&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
-                </a>
-            </div>
-            """
-        )
-        
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt (long prompts are automatically supported)",
-                container=False,
-                value=DEFAULT_PROMPT  # Set default prompt
-            )
-
-            run_button = gr.Button("Run", scale=0)
-
-        result = gr.Image(format="png", label="Result", show_label=False)
-        
-        with gr.Accordion("Advanced Settings", open=False):
-
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-#                value="bad quality,worst quality,worst detail,sketch,censor,"
-                value="monochrome, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn,"
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=MAX_IMAGE_SIZE,
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=20.0,
-                    step=0.1,
-                    value=7,
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=28,
-                    step=1,
-                    value=28,
-                )
-
-    run_button.click(
-        fn=infer,
-        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs=[result]
-    )
-
-demo.queue().launch()
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+        error_img = Image.new('RGB', (width, height), color=(100, 0, 0))
+        return error_img