Update app.py

app.py

@@ -2,7 +2,7 @@ import spaces
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM, pipeline, Qwen2VLForConditionalGeneration
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, pipeline
 from diffusers import DiffusionPipeline
 import random
 import numpy as np
@@ -20,7 +20,9 @@ dtype = torch.bfloat16
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
 # FLUX.1-dev model
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token=huggingface_token).to(device)
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token=huggingface_token
+).to(device)
 
 # Initialize Qwen2VL model
 qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
@@ -32,7 +34,7 @@ qwen_processor = AutoProcessor.from_pretrained("prithivMLmods/JSONify-Flux", tru
 enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
+MAX_IMAGE_SIZE = 1024  # Reduced to prevent memory issues
 
 # Qwen2VL caption function
 @spaces.GPU
@@ -101,14 +103,24 @@ def process_workflow(image, text_prompt, use_enhancer, seed, randomize_seed, wid
     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    image = pipe(
-        prompt=prompt,
-        generator=generator,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale
-    ).images[0]
+    # Reduce memory usage by clearing GPU cache
+    torch.cuda.empty_cache()
+    
+    # Generate image with FLUX.1-dev
+    try:
+        image = pipe(
+            prompt=prompt,
+            generator=generator,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale
+        ).images[0]
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            raise RuntimeError("CUDA out of memory. Try reducing image size or inference steps.")
+        else:
+            raise e
     
     return image, prompt, seed
 
@@ -151,10 +163,10 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
                 use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False)
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                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=1024)
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)  # Reduced default width
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512)  # Reduced default height
                 guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=15, step=0.1, value=3.5)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=28)
+                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=20)  # Reduced default steps
             
             generate_btn = gr.Button("Generate Image", elem_classes="submit-btn")