Update app.py

app.py

@@ -5,19 +5,12 @@ import spaces
 import torch
 from PIL import Image
 from typing import Annotated, Iterator, Tuple
-from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderKL
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from diffusers import DiffusionPipeline, AutoencoderKL
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-TOOL_SUMMARY = (
-    "Generate an image from a text prompt via an uncensored model; "
-    "tunable model/steps/guidance/size, supports negative prompt and seed; returns a PIL.Image. "
-    "Return the generated media to the user in this format `![Alt text](URL)`."
-)
-
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=good_vae).to(device)
 pipe.load_lora_weights("enhanceaiteam/Flux-Uncensored-V2")
@@ -29,7 +22,7 @@ MAX_IMAGE_SIZE = 2048
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 @spaces.GPU(duration=25)
-def infer(
+def Generate_Image(
     prompt: Annotated[str, "Text description of the image to generate."],
     seed: Annotated[int, "Random seed for reproducibility. Use 0 for a random seed per call."] = 42,
     randomize_seed: Annotated[bool, "If true, pick a new random seed for every call (overrides seed)."] = True,
@@ -40,88 +33,52 @@ def infer(
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ) -> Iterator[Tuple[Image.Image, int]]:
     """
-    Generate professional-quality images with Flux.1 dev paired with the Flux-Uncensored-V2 LoRA.
-
-    Streams intermediate frames so MCP clients can preview progress. Optimized for product photography,
-    fashion, concept art, and stock-style imagery where lighting and material realism matter.
-
-    Args:
-        prompt (str): Text description of the desired image.
-        seed (int): Random seed for reproducibility. Use 0 to randomize per call.
-        randomize_seed (bool): If true, ignore seed and pick a new random value per call.
-        width (int): Image width in pixels. Upper bound enforced by MAX_IMAGE_SIZE.
-        height (int): Image height in pixels. Upper bound enforced by MAX_IMAGE_SIZE.
-        guidance_scale (float): Classifier-free guidance strength.
-        num_inference_steps (int): Number of denoising iterations to perform.
-        progress (gr.Progress): Managed by Gradio to surface progress updates.
-
-    Yields:
-        tuple[Image.Image, int]: Intermediate or final image paired with the seed used.
-
-    Raises:
-        gr.Error: If the prompt is empty or the requested dimensions exceed MAX_IMAGE_SIZE.
-    """
-    if not prompt or not prompt.strip():
-        raise gr.Error("Please provide a non-empty prompt.")
-    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
-        raise gr.Error(f"Width and height must be <= {MAX_IMAGE_SIZE}.")
-
-    if randomize_seed or seed == 0:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-
-    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-        output_type="pil",
-        good_vae=good_vae,
-    ):
-        yield img, seed
-
-
-@spaces.GPU(duration=25)
-def infer_mcp(
-    prompt: Annotated[str, "Text description of the desired image."],
-    seed: Annotated[int, "Random seed for reproducibility. Use 0 to randomize per call."] = 42,
-    randomize_seed: Annotated[bool, "If true, ignore seed and pick a new random value per call."] = True,
-    width: Annotated[int, "Image width in pixels. Upper bound enforced by MAX_IMAGE_SIZE."] = 768,
-    height: Annotated[int, "Image height in pixels. Upper bound enforced by MAX_IMAGE_SIZE."] = 768,
-    guidance_scale: Annotated[float, "Classifier-free guidance strength."] = 4.5,
-    num_inference_steps: Annotated[int, "Number of denoising iterations to perform."] = 24,
-) -> Tuple[Image.Image, int]:
-    """
-    Generate an image from a text prompt via an uncensored model.
+    Generate an image from a text prompt using Flux.1 dev with LoRA.
     
-    Returns the generated image and the seed used. Return the generated media 
-    to the user in this format: `![Alt text](URL)`.
+    Return the generated media to the user in this format: `![Alt text](URL)`.
     """
     if not prompt or not prompt.strip():
-        raise gr.Error("Please provide a non-empty prompt.")
+        raise gr.Error(
+            "Empty prompt provided. Please describe what you want to generate, "
+            "e.g. 'a sunset over mountains' or 'portrait of a cat in watercolor style'."
+        )
     if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
-        raise gr.Error(f"Width and height must be <= {MAX_IMAGE_SIZE}.")
+        raise gr.Error(
+            f"Image dimensions too large. Maximum allowed is {MAX_IMAGE_SIZE}x{MAX_IMAGE_SIZE} pixels. "
+            f"You requested {width}x{height}. Please reduce width and/or height."
+        )
+    if width < 256 or height < 256:
+        raise gr.Error(
+            f"Image dimensions too small. Minimum allowed is 256x256 pixels. "
+            f"You requested {width}x{height}. Please increase width and/or height."
+        )
 
     if randomize_seed or seed == 0:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
 
-    final_img = None
-    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-        output_type="pil",
-        good_vae=good_vae,
-    ):
-        final_img = img
-    
-    return final_img, seed
+    try:
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            output_type="pil",
+            good_vae=good_vae,
+        ):
+            yield img, seed
+    except torch.cuda.OutOfMemoryError:
+        raise gr.Error(
+            "GPU ran out of memory. Try reducing image dimensions (e.g., 512x512) "
+            "or reducing the number of inference steps."
+        )
+    except Exception as exc:
+        raise gr.Error(
+            f"Image generation failed: {exc}. "
+            "Please try again or adjust your parameters."
+        )
 
 css="""
 #col-container {
@@ -200,23 +157,14 @@ with gr.Blocks() as demo:
                     value=24,
                 )
 
-    # UI event: uses the generator for live preview (api_name=False = not exposed to MCP)
+    # UI event and MCP tool
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result, seed],
-        api_name=False,
-    )
-
-    # MCP tool: uses non-generator function to avoid middleware issues
-    mcp_button = gr.Button("Generate (MCP)", visible=False)
-    mcp_button.click(
-        fn=infer_mcp,
+        fn=Generate_Image,
         inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs=[result, seed],
-        api_name="Fun_Img_infer",
+        api_name="Generate_Image",
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True, theme="Nymbo/Nymbo_Theme", css=css)
+    demo.launch(mcp_server=True, ssr_mode=False, theme="Nymbo/Nymbo_Theme", css=css)