app.py

import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# Apply more comprehensive patches to Gradio's utility functions
import gradio_client.utils
import types

# Patch 1: Fix the _json_schema_to_python_type function
original_json_schema = gradio_client.utils._json_schema_to_python_type

def patched_json_schema(schema, defs=None):
    # Handle boolean values directly
    if isinstance(schema, bool):
        return "bool"
    
    # Handle cases where 'additionalProperties' is a boolean
    try:
        if "additionalProperties" in schema and isinstance(schema["additionalProperties"], bool):
            schema["additionalProperties"] = {"type": "any"}
    except (TypeError, KeyError):
        pass
    
    # Call the original function
    try:
        return original_json_schema(schema, defs)
    except Exception as e:
        # Fallback to a safe value when the schema can't be parsed
        return "any"

# Replace the original function with our patched version
gradio_client.utils._json_schema_to_python_type = patched_json_schema

# Create permanent storage directory
SAVE_DIR = "saved_images"  # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

# Safe settings for model loading
device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "openfree/flux-chatgpt-ghibli-lora"

def load_model_with_retry(max_retries=5):
    for attempt in range(max_retries):
        try:
            print(f"Loading model attempt {attempt+1}/{max_retries}...")
            pipeline = DiffusionPipeline.from_pretrained(
                repo_id, 
                torch_dtype=torch.bfloat16,
                use_safetensors=True,
                resume_download=True
            )
            print("Model loaded successfully, loading LoRA weights...")
            pipeline.load_lora_weights(adapter_id)
            pipeline = pipeline.to(device)
            print("Pipeline ready!")
            return pipeline
        except Exception as e:
            if attempt < max_retries - 1:
                wait_time = 10 * (attempt + 1)
                print(f"Error loading model: {e}. Retrying in {wait_time} seconds...")
                import time
                time.sleep(wait_time)
            else:
                raise Exception(f"Failed to load model after {max_retries} attempts: {e}")

# Load the model
pipeline = load_model_with_retry()

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def save_generated_image(image, prompt):
    # Generate unique filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    filename = f"{timestamp}_{unique_id}.png"
    filepath = os.path.join(SAVE_DIR, filename)
    
    # Save the image
    image.save(filepath)
    
    # Save metadata
    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
    with open(metadata_file, "a", encoding="utf-8") as f:
        f.write(f"{filename}|{prompt}|{timestamp}\n")
    
    return filepath

def load_generated_images():
    if not os.path.exists(SAVE_DIR):
        return []
    
    # Load all images from the directory
    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    # Sort by creation time (newest first)
    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
    return image_files

@spaces.GPU(duration=120)
def inference(
    prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
    lora_scale: float,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # Error handling for the inference process
    try:
        image = pipeline(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            joint_attention_kwargs={"scale": lora_scale},
        ).images[0]
        
        # Save the generated image
        filepath = save_generated_image(image, prompt)
        
        # Return the image, seed, and updated gallery
        return image, seed, load_generated_images()
    except Exception as e:
        # Log the error and return a simple error image
        print(f"Error during inference: {e}")
        error_img = Image.new('RGB', (width, height), color='red')
        return error_img, seed, load_generated_images()

examples = [
    "Ghibli style futuristic stormtrooper with glossy white armor and a sleek helmet, standing heroically on a lush alien planet, vibrant flowers blooming around, soft sunlight illuminating the scene, a gentle breeze rustling the leaves. The armor reflects the pink and purple hues of the alien sunset, creating an ethereal glow around the figure. [trigger]",
    
    "Ghibli style young mechanic girl in a floating workshop, surrounded by hovering tools and glowing mechanical parts, her blue overalls covered in oil stains, tinkering with a semi-transparent robot companion. Magical sparks fly as she works, while floating islands with waterfalls drift past her open workshop window. [trigger]",
    
    "Ghibli style ancient forest guardian robot, covered in moss and flowering vines, sitting peacefully in a crystal-clear lake. Its gentle eyes glow with soft blue light, while bioluminescent dragonflies dance around its weathered metal frame. Ancient tech symbols on its surface pulse with a gentle rhythm. [trigger]",
    
    "Ghibli style sky whale transport ship, its metallic skin adorned with traditional Japanese patterns, gliding through cotton candy clouds at sunrise. Small floating gardens hang from its sides, where workers in futuristic kimonos tend to glowing plants. Rainbow auroras shimmer in the background. [trigger]",
    
    "Ghibli style cyber-shrine maiden with flowing holographic robes, performing a ritual dance among floating lanterns and digital cherry blossoms. Her traditional headdress emits soft light patterns, while spirit-like AI constructs swirl around her in elegant patterns. The scene is set in a modern shrine with both ancient wood and sleek chrome elements. [trigger]",
    
    "Ghibli style robot farmer tending to floating rice paddies in the sky, wearing a traditional straw hat with advanced sensors. Its gentle movements create ripples in the water as it plants glowing rice seedlings. Flying fish leap between the terraced fields, leaving trails of sparkles in their wake, while future Tokyo's spires gleam in the distance. [trigger]"
]

css = """
footer {
    visibility: hidden;
}
"""

# Use a simpler UI configuration that is less likely to cause issues
with gr.Blocks(css=css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title"> FLUX Ghibli LoRA</div>')

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt")
            
            with gr.Row():
                run_button = gr.Button("Generate Image")
                clear_button = gr.Button("Clear")
            
            with gr.Accordion("Settings", open=False):
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42,
                )
                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=768,
                    )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=3.5,
                    )
                    num_inference_steps = gr.Slider(
                        label="Steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=30,
                    )
                    lora_scale = gr.Slider(
                        label="LoRA scale",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.1,
                        value=1.0,
                    )
            
            gr.Examples(
                examples=examples,
                inputs=prompt,
            )
        
        with gr.Column(scale=4):
            result = gr.Image(label="Generated Image")
            seed_text = gr.Number(label="Used Seed", value=42)
    
    with gr.Tab("Gallery"):
        gallery_header = gr.Markdown("### Generated Images Gallery")
        generated_gallery = gr.Gallery(
            label="Generated Images",
            columns=3,
            value=load_generated_images(),
            height="auto"
        )
        refresh_btn = gr.Button("🔄 Refresh Gallery")

    # Event handlers
    def refresh_gallery():
        return load_generated_images()

    def clear_output():
        return "", gr.update(value=None), seed

    refresh_btn.click(
        fn=refresh_gallery,
        inputs=None,
        outputs=generated_gallery,
    )
    
    clear_button.click(
        fn=clear_output,
        inputs=None,
        outputs=[prompt, result, seed_text]
    )

    run_button.click(
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed_text, generated_gallery],
    )
    
    prompt.submit(
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed_text, generated_gallery],
    )

# Launch with fallback options
try:
    demo.queue(concurrency_count=1, max_size=10)
    demo.launch(debug=True, show_api=False)
except Exception as e:
    print(f"Error during launch: {e}")
    print("Trying alternative launch configuration...")
    # Skip queue and simplify launch parameters
    demo.launch(debug=True, show_api=False, share=False)