Update app.py

app.py

@@ -6,9 +6,7 @@ import random
 import spaces
 import torch
 from diffusers import Flux2Pipeline, Flux2Transformer2DModel
-import requests
 from PIL import Image
-import base64
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -16,45 +14,19 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def remote_text_encoder(prompts, max_retries=3):
-    from gradio_client import Client
-    import time
-    
-    for attempt in range(max_retries):
-        try:
-            client = Client("multimodalart/mistral-text-encoder")
-            result = client.predict(
-                prompt=prompts,
-                api_name="/encode_text"
-            )
-            prompt_embeds = torch.load(result[0])
-            return prompt_embeds
-        except Exception as e:
-            print(f"Text encoder attempt {attempt + 1}/{max_retries} failed: {e}")
-            if attempt < max_retries - 1:
-                time.sleep(2)
-            else:
-                raise Exception(f"Text encoder failed after {max_retries} attempts: {e}")
-
-# Load model
+# Load the full pipeline WITH text encoder
+print("Loading FLUX.2 pipeline with text encoder...")
 repo_id = "black-forest-labs/FLUX.2-dev"
 
-dit = Flux2Transformer2DModel.from_pretrained(
-    repo_id,
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16
-)
-
 pipe = Flux2Pipeline.from_pretrained(
     repo_id,
-    text_encoder=None,
-    transformer=dit,
     torch_dtype=torch.bfloat16
 )
-pipe.to(device)
 
-# AOTI blocks temporarily disabled - HuggingFace needs to recompile for new ZeroGPU environment
-# spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
+# Enable memory optimizations for ZeroGPU
+pipe.enable_model_cpu_offload()  # This offloads models to CPU when not in use, saving VRAM
+
+print("Pipeline loaded successfully!")
 
 def update_dimensions_from_image(image_list):
     """Update width/height sliders based on uploaded image aspect ratio."""
@@ -81,33 +53,7 @@ def update_dimensions_from_image(image_list):
     
     return new_width, new_height
 
-def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
-    num_images = 0 if image_list is None else len(image_list)
-    step_duration = 1 + 0.8 * num_images
-    return max(45, num_inference_steps * step_duration + 10)
-
-@spaces.GPU(duration=get_duration)
-def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
-    prompt_embeds = prompt_embeds.to(device)
-    
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    pipe_kwargs = {
-        "prompt_embeds": prompt_embeds,
-        "image": image_list,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-        "width": width,
-        "height": height,
-    }
-    
-    if progress:
-        progress(0, desc="Starting generation...")
-        
-    image = pipe(**pipe_kwargs).images[0]
-    return image
-
+@spaces.GPU(duration=120)  # Increased duration since we're doing text encoding + generation
 def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
@@ -119,30 +65,34 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         for item in input_images:
             image_list.append(item[0])
 
-    # Text Encoding
-    progress(0.1, desc="Encoding prompt...")
-    prompt_embeds = remote_text_encoder(prompt)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     
-    # Image Generation
-    progress(0.3, desc="Waiting for GPU...")
-    image = generate_image(
-        prompt_embeds, 
-        image_list, 
-        width, 
-        height, 
-        num_inference_steps, 
-        guidance_scale, 
-        seed, 
-        progress
-    )
+    progress(0.1, desc="Encoding prompt & generating image...")
+    
+    # Build pipeline arguments
+    pipe_kwargs = {
+        "prompt": prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+        "width": width,
+        "height": height,
+    }
+    
+    # Add images if provided
+    if image_list is not None:
+        pipe_kwargs["image"] = image_list
+    
+    # Run the pipeline - text encoding happens automatically inside
+    image = pipe(**pipe_kwargs).images[0]
     
     return image, seed
 
 examples = [
     ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
-    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
+    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire."],
     ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
-    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
+    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose."],
 ]
 
 examples_images = [
@@ -159,7 +109,7 @@ css="""
 }
 """
 
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# FLUX.2 [dev]
@@ -240,8 +190,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
             fn=infer,
             inputs=[prompt],
             outputs=[result, seed],
-            cache_examples=True,
-            cache_mode="lazy"
+            cache_examples="lazy"
         )
 
         gr.Examples(
@@ -249,8 +198,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
             fn=infer,
             inputs=[prompt, input_images],
             outputs=[result, seed],
-            cache_examples=True,
-            cache_mode="lazy"
+            cache_examples="lazy"
         )
 
     input_images.upload(
@@ -266,4 +214,4 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
         outputs=[result, seed]
     )
 
-demo.launch(css=css)
+demo.launch()