Update app.py

app.py

@@ -9,8 +9,8 @@ from diffusers import DiffusionPipeline
 from custom_pipeline import FLUXPipelineWithIntermediateOutputs
 from transformers import pipeline
 
-# Translation model loading
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+# Translation model loading with device specification
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -18,13 +18,19 @@ MAX_IMAGE_SIZE = 2048
 DEFAULT_WIDTH = 1024
 DEFAULT_HEIGHT = 1024
 DEFAULT_INFERENCE_STEPS = 1
+GPU_DURATION = 15  # Reduced from 25 to stay within quota
+
+# Device and model setup with memory optimization
+def setup_model():
+    dtype = torch.float16
+    pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell", 
+        torch_dtype=dtype,
+        device_map="auto"  # Enable model parallelism
+    )
+    return pipe
 
-# Device and model setup
-dtype = torch.float16
-pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
-).to("cuda")
-torch.cuda.empty_cache()
+pipe = setup_model()
 
 # Menu labels dictionary
 english_labels = {
@@ -41,36 +47,67 @@ english_labels = {
 }
 
 def translate_if_korean(text):
-    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
-        return translator(text)[0]['translation_text']
-    return text
-
-# Modified inference function to always use random seed for examples
-@spaces.GPU(duration=25)
-def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=True, num_inference_steps=DEFAULT_INFERENCE_STEPS):
-    prompt = translate_if_korean(prompt)
-    
-    # Always generate a random seed if none provided or randomize_seed is True
-    if seed is None or randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-
-    start_time = time.time()
-
-    for img in pipe.generate_images(  
-            prompt=prompt,
-            guidance_scale=0,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator
-        ): 
-        latency = f"Processing Time: {(time.time()-start_time):.2f} seconds"    
-        yield img, seed, latency
-
-# Function specifically for examples that always uses random seeds
+    """Safely translate Korean text to English."""
+    try:
+        if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
+            return translator(text)[0]['translation_text']
+        return text
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return text
+
+# Modified inference function with error handling and memory management
+@spaces.GPU(duration=GPU_DURATION)
+def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, 
+                  randomize_seed=True, num_inference_steps=DEFAULT_INFERENCE_STEPS):
+    try:
+        # Input validation
+        if not isinstance(seed, (int, type(None))):
+            seed = None
+            randomize_seed = True
+            
+        prompt = translate_if_korean(prompt)
+        
+        if seed is None or randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+        # Ensure valid dimensions
+        width = min(max(256, width), MAX_IMAGE_SIZE)
+        height = min(max(256, height), MAX_IMAGE_SIZE)
+        
+        generator = torch.Generator().manual_seed(seed)
+        
+        start_time = time.time()
+        
+        with torch.cuda.amp.autocast():  # Enable automatic mixed precision
+            for img in pipe.generate_images(
+                prompt=prompt,
+                guidance_scale=0,
+                num_inference_steps=num_inference_steps,
+                width=width,
+                height=height,
+                generator=generator
+            ):
+                latency = f"Processing Time: {(time.time()-start_time):.2f} seconds"
+                
+                # Clear CUDA cache after generation
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                    
+                yield img, seed, latency
+                
+    except Exception as e:
+        print(f"Error in generate_image: {e}")
+        # Return a blank image or error message
+        yield None, seed, f"Error: {str(e)}"
+
+# Example generator with error handling
 def generate_example_image(prompt):
-    return generate_image(prompt, randomize_seed=True)
+    try:
+        return next(generate_image(prompt, randomize_seed=True))
+    except Exception as e:
+        print(f"Error in example generation: {e}")
+        return None, None, f"Error: {str(e)}"
 
 # Example prompts
 examples = [
@@ -88,12 +125,14 @@ footer {
 }
 """
 
-# --- Gradio UI ---
+# --- Gradio UI with improved error handling ---
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     with gr.Column(elem_id="app-container"):
         with gr.Row():
             with gr.Column(scale=3):
-                result = gr.Image(label=english_labels["Generated Image"], show_label=False, interactive=False)
+                result = gr.Image(label=english_labels["Generated Image"], 
+                                show_label=False, 
+                                interactive=False)
             with gr.Column(scale=1):
                 prompt = gr.Text(
                     label=english_labels["Prompt"],
@@ -108,25 +147,53 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                     with gr.Row():
                         latency = gr.Text(show_label=False)
                     with gr.Row():
-                        seed = gr.Number(label=english_labels["Seed"], value=42, precision=0)
-                        randomize_seed = gr.Checkbox(label=english_labels["Randomize Seed"], value=True)
+                        # Modified Number component with proper validation
+                        seed = gr.Number(
+                            label=english_labels["Seed"], 
+                            value=42, 
+                            precision=0,
+                            minimum=0,
+                            maximum=MAX_SEED
+                        )
+                        randomize_seed = gr.Checkbox(
+                            label=english_labels["Randomize Seed"], 
+                            value=True
+                        )
                     with gr.Row():
-                        width = gr.Slider(label=english_labels["Width"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
-                        height = gr.Slider(label=english_labels["Height"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
-                        num_inference_steps = gr.Slider(label=english_labels["Inference Steps"], minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)
+                        width = gr.Slider(
+                            label=english_labels["Width"], 
+                            minimum=256, 
+                            maximum=MAX_IMAGE_SIZE, 
+                            step=32, 
+                            value=DEFAULT_WIDTH
+                        )
+                        height = gr.Slider(
+                            label=english_labels["Height"], 
+                            minimum=256, 
+                            maximum=MAX_IMAGE_SIZE, 
+                            step=32, 
+                            value=DEFAULT_HEIGHT
+                        )
+                        num_inference_steps = gr.Slider(
+                            label=english_labels["Inference Steps"], 
+                            minimum=1, 
+                            maximum=4, 
+                            step=1, 
+                            value=DEFAULT_INFERENCE_STEPS
+                        )
 
         with gr.Row():
             gr.Markdown(f"### 🌟 {english_labels['Inspiration Gallery']}")
         with gr.Row():
             gr.Examples(
                 examples=examples,
-                fn=generate_example_image,  # Use the example-specific function
+                fn=generate_example_image,
                 inputs=[prompt],
                 outputs=[result, seed],
-                cache_examples=False  # Disable caching to ensure new generation each time
+                cache_examples=False
             )
 
-    # Event handling
+    # Event handling with improved error handling
     enhanceBtn.click(
         fn=generate_image,
         inputs=[prompt, seed, width, height],
@@ -136,9 +203,17 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         queue=False
     )
 
+    # Modified event handler with proper input validation
+    def validated_generate(*args):
+        try:
+            return next(generate_image(*args))
+        except Exception as e:
+            print(f"Error in validated_generate: {e}")
+            return None, args[1], f"Error: {str(e)}"
+
     gr.on(
         triggers=[prompt.input, width.input, height.input, num_inference_steps.input],
-        fn=generate_image,
+        fn=validated_generate,
         inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
         outputs=[result, seed, latency],
         show_progress="hidden",
@@ -147,4 +222,5 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         queue=False
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()