Update app.py

app.py

@@ -9,18 +9,71 @@ 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"
 
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-pipeline.load_lora_weights(adapter_id)
-pipeline = pipeline.to(device)
+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
@@ -53,10 +106,6 @@ def load_generated_images():
     image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
     return image_files
 
-def load_predefined_images():
-    # Return empty list since we're not using predefined images
-    return []
-
 @spaces.GPU(duration=120)
 def inference(
     prompt: str,
@@ -73,21 +122,28 @@ def inference(
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    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()
+    # 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]",
@@ -109,105 +165,121 @@ footer {
 }
 """
 
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
+# 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>')
-    gr.HTML('<div class="title">😄Image to Video Explore: <a href="https://huggingface.co/spaces/ginigen/theater" target="_blank">https://huggingface.co/spaces/ginigen/theater</a></div>')
 
-    with gr.Tabs() as tabs:
-        with gr.Tab("Generation"):
-            with gr.Column(elem_id="col-container"):
+    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():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt",
-                        container=False,
+                    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,
                     )
-                    run_button = gr.Button("Run", scale=0)
-
-                result = gr.Image(label="Result", show_label=False)
 
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
+                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=42,
+                        value=30,
                     )
-                    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="Number of inference 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],
-                    outputs=[result, seed],
-                )
-
-        with gr.Tab("Gallery"):
-            gallery_header = gr.Markdown("### Generated Images Gallery")
-            generated_gallery = gr.Gallery(
-                label="Generated Images",
-                columns=6,
-                show_label=False,
-                value=load_generated_images(),
-                elem_id="generated_gallery",
-                height="auto"
+                    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,
             )
-            refresh_btn = gr.Button("🔄 Refresh Gallery")
-
+        
+        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]
+    )
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
+    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,
@@ -219,8 +291,15 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=Fa
             num_inference_steps,
             lora_scale,
         ],
-        outputs=[result, seed, generated_gallery],
+        outputs=[result, seed_text, generated_gallery],
     )
 
-demo.queue()
-demo.launch()
+# 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)