Update app.py

app.py

@@ -1,22 +1,23 @@
 import os
 import subprocess
 import sys
+
+subprocess.check_call(["unzip", "-o", "diffusers.zip", "-d", "diffusers_local"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "diffusers_local/diffusers"])
+
 import io
 import gradio as gr
 import numpy as np
 import random
 import spaces
 import torch
-from diffusers import Flux2Pipeline, Flux2Transformer2DModel
-from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
+from diffusers import Flux2KleinPipeline
 import requests
 from PIL import Image
 import json
 import base64
 from huggingface_hub import InferenceClient
 
-subprocess.check_call([sys.executable, "-m", "pip", "install", "spaces==0.43.0"])
-
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -50,38 +51,13 @@ Rules:
 
 Output only the final instruction in plain text and nothing else."""
 
-def remote_text_encoder(prompts):
-    from gradio_client import Client
-    
-    client = Client("multimodalart/mistral-text-encoder")
-    result = client.predict(
-        prompt=prompts,
-        api_name="/encode_text"
-    )
-    
-    # Load returns a tensor, usually on CPU by default
-    prompt_embeds = torch.load(result[0])
-    return prompt_embeds
-
-# Load model
-repo_id = "black-forest-labs/FLUX.2-dev"
+# Load model - using Flux2KleinPipeline with built-in text encoder
+repo_id = "diffusers-internal-dev/dummy-1015-4b"  # 4b model
+# repo_id = "diffusers-internal-dev/dummy-1015-9b"  # 9b model (alternative)
 
-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)
+pipe = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
+pipe.enable_model_cpu_offload()
 
-# Pull pre-compiled Flux2 Transformer blocks from HF hub
-spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
 
 def image_to_data_uri(img):
     buffered = io.BytesIO()
@@ -89,6 +65,7 @@ def image_to_data_uri(img):
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
     return f"data:image/png;base64,{img_str}"
 
+
 def upsample_prompt_logic(prompt, image_list):
     try:
         if image_list and len(image_list) > 0:
@@ -128,6 +105,7 @@ def upsample_prompt_logic(prompt, image_list):
         print(f"Upsampling failed: {e}")
         return prompt
 
+
 def update_dimensions_from_image(image_list):
     """Update width/height sliders based on uploaded image aspect ratio.
     Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
@@ -157,37 +135,9 @@ def update_dimensions_from_image(image_list):
     
     return new_width, new_height
 
-# Updated duration function to match generate_image arguments (including progress)
-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(65, 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)):
-    # Move embeddings to GPU only when inside the GPU decorated function
-    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,
-    }
-    
-    # Progress bar for the actual generation steps
-    if progress:
-        progress(0, desc="Starting generation...")
-        
-    image = pipe(**pipe_kwargs).images[0]
-    return image
 
-def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU
+def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -199,34 +149,37 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         for item in input_images:
             image_list.append(item[0])
 
-    # 1. Upsampling (Network bound - No GPU needed)
+    # 1. Upsampling (Network bound)
     final_prompt = prompt
     if prompt_upsampling:
-        progress(0.05, desc="Upsampling prompt...")
+        progress(0.1, desc="Upsampling prompt...")
         final_prompt = upsample_prompt_logic(prompt, image_list)
         print(f"Original Prompt: {prompt}")
         print(f"Upsampled Prompt: {final_prompt}")
 
-    # 2. Text Encoding (Network bound - No GPU needed)
-    progress(0.1, desc="Encoding prompt...")
-    # This returns CPU tensors
-    prompt_embeds = remote_text_encoder(final_prompt)
+    # 2. Image Generation
+    progress(0.2, desc="Generating image...")
     
-    # 3. Image Generation (GPU bound)
-    progress(0.3, desc="Waiting for GPU...")
-    image = generate_image(
-        prompt_embeds, 
-        image_list, 
-        width, 
-        height, 
-        num_inference_steps, 
-        guidance_scale, 
-        seed, 
-        progress
-    )
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    pipe_kwargs = {
+        "prompt": final_prompt,
+        "height": height,
+        "width": width,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+    }
+    
+    # Add images if provided
+    if image_list is not None:
+        pipe_kwargs["image"] = image_list
+    
+    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."],
@@ -235,11 +188,10 @@ examples = [
 ]
 
 examples_images = [
-    # ["Replace the top of the person from image 1 with the one from image 2", ["person1.webp", "woman2.webp"]],
     ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
 ]
 
-css="""
+css = """
 #col-container {
     margin: 0 auto;
     max-width: 1200px;
@@ -249,11 +201,11 @@ 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]
-FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and combining images based on text instructions model [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
+        gr.Markdown(f"""# FLUX.2 Klein [dev]
+FLUX.2 Klein [dev] is a distilled model capable of generating, editing and combining images based on text instructions [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
         """)
         with gr.Row():
             with gr.Column():
@@ -319,7 +271,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                             minimum=1,
                             maximum=100,
                             step=1,
-                            value=30,
+                            value=50,
                         )
                         
                         guidance_scale = gr.Slider(
@@ -327,7 +279,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                             minimum=0.0,
                             maximum=10.0,
                             step=0.1,
-                            value=4,
+                            value=4.0,
                         )
                 
                 
@@ -367,4 +319,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()