Squashing commit

Co-authored-by: linoyts <linoyts@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+bird.webp filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,12 @@
+---
+title: FLUX.2-dev
+emoji: 💻
+colorFrom: red
+colorTo: gray
+sdk: gradio
+sdk_version: 6.0.0
+app_file: app.py
+pinned: true
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,256 @@
+import os
+import subprocess
+import sys
+import io
+from kernels import get_kernel
+
+# Clone and install diffusers from private repo
+def setup_diffusers():
+    gh_token = os.environ.get("GH_API_KEY")
+    if not gh_token:
+        raise ValueError("GH_API_KEY environment variable not set")
+    
+    repo_url = f"https://{gh_token}@github.com/apolinario/diffusers-new-model-addition.git"
+    repo_dir = "diffusers"
+    
+    # Clone repo if it doesn't exist
+    if not os.path.exists(repo_dir):
+        print("Cloning diffusers repository...")
+        subprocess.run(["git", "clone", repo_url, repo_dir], check=True)
+    
+    # Checkout the branch
+    #print("Checking out full-flux2 branch...")
+    #subprocess.run(["git", "-C", repo_dir, "checkout", "full-flux2"], check=True)
+        
+    # Install in editable mode
+    print("Installing diffusers in editable mode...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "-e", repo_dir], check=True)
+
+    # Add to Python path
+    diffusers_src = os.path.abspath(os.path.join(repo_dir, "src"))
+    if diffusers_src not in sys.path:
+        sys.path.insert(0, diffusers_src)
+    
+    print("Diffusers setup complete!")
+
+# Run setup
+setup_diffusers()
+
+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 optimization import optimize_pipeline_
+import requests
+from PIL import Image
+import json
+
+dtype = torch.bfloat16
+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):
+    response = requests.post(
+        "https://remote-text-encoder-flux-2.huggingface.co/predict",
+        json={"prompt": prompts},
+        headers={
+            "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
+            "Content-Type": "application/json"
+        }
+    )
+
+    assert response.status_code == 200, f"{response.status_code=}"
+
+    prompt_embeds = torch.load(io.BytesIO(response.content))
+
+    return prompt_embeds
+
+# Load model
+fa3_kernel = get_kernel("kernels-community/flash-attn3", revision="fake-ops-return-probs")
+
+repo_id = "diffusers-internal-dev/new-model-image-final-weights"
+dit = Flux2Transformer2DModel.from_pretrained(
+    repo_id,
+    subfolder="transformer",
+    attn_implementation=fa3_kernel,
+    torch_dtype=torch.bfloat16
+)
+
+pipe = Flux2Pipeline.from_pretrained(
+    repo_id,
+    text_encoder=None,
+    transformer=dit,
+    torch_dtype=torch.bfloat16
+)
+pipe.to("cuda")
+# pipe.transformer.compile_repeated_blocks(dynamic=True)
+
+#optimize_pipeline_(
+#    pipe,
+#    prompt_embeds=remote_text_encoder("prompt").to("cuda"),
+#    guidance_scale=2.5,
+#    width=1024,
+#    height=1024,
+#    num_inference_steps=1,
+#)
+
+@spaces.GPU(duration=180)
+def infer(prompt, input_images, 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:
+        seed = random.randint(0, MAX_SEED)
+    
+    # Get prompt embeddings from remote text encoder
+    progress(0.1, desc="Encoding prompt...")
+    prompt_embeds = remote_text_encoder(prompt).to("cuda")
+    
+    # Prepare image list (convert None or empty gallery to None)
+    image_list = None
+    if input_images is not None and len(input_images) > 0:
+        image_list = []
+        for item in input_images:
+            image_list.append(item[0])
+    
+    # Generate image
+    progress(0.3, desc="Generating image...")
+    generator = torch.Generator(device=device).manual_seed(seed)
+    image = pipe(
+        prompt_embeds=prompt_embeds,
+        image=image_list,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+    ).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", None],
+    ["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.", None],
+    ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo", None],
+    ["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.", None],
+]
+
+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="""
+#col-container {
+    margin: 0 auto;
+    max-width: 620px;
+}
+"""
+
+with gr.Blocks() as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# FLUX.2 [dev]
+32B param rectified flow transformer guidance-distilled from [FLUX.2 pro](#) [[model](#)], [[blog](#)]
+        """)
+
+        with gr.Accordion("Input image(s) (optional)", open=False):
+            input_images = gr.Gallery(
+                label="Input Image(s)",
+                type="pil",
+                columns=3,
+                rows=1,
+            )
+        
+        with gr.Row():
+            
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=2,
+                placeholder="Enter your prompt",
+                container=False,
+                scale=3
+            )
+            
+            run_button = gr.Button("Run", scale=1)
+        
+        result = gr.Image(label="Result", show_label=False)
+        
+        with gr.Accordion("Advanced Settings", open=False):
+            
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            
+            with gr.Row():
+                
+                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,
+                )
+            
+            with gr.Row():
+                
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=30,
+                )
+                
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=4,
+                )
+        
+        gr.Examples(
+            examples=examples,
+            fn=infer,
+            inputs=[prompt, input_images],
+            outputs=[result, seed],
+            cache_examples=True,
+            cache_mode="lazy"
+        )
+
+        gr.Examples(
+            examples=examples_images,
+            fn=infer,
+            inputs=[prompt, input_images],
+            outputs=[result, seed],
+            cache_examples=True,
+            cache_mode="lazy"
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
+        outputs=[result, seed]
+    )
+
+demo.launch(css=css)

optimization.py

@@ -0,0 +1,60 @@
+"""
+"""
+
+from typing import Any
+from typing import Callable
+from typing import ParamSpec
+import spaces
+import torch
+from torch.utils._pytree import tree_map
+
+P = ParamSpec('P')
+
+TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length', min=64, max=16384)
+TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length', min=64, max=512)
+
+TRANSFORMER_DYNAMIC_SHAPES = {
+    'hidden_states': {
+        1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
+    },
+    'encoder_hidden_states': {
+        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
+    },
+    'img_ids': {
+        1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
+    },
+    'txt_ids': {
+        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM, 
+    },
+}
+
+INDUCTOR_CONFIGS = {
+    'conv_1x1_as_mm': True,
+    'epilogue_fusion': False,
+    'coordinate_descent_tuning': True,
+    'coordinate_descent_check_all_directions': True,
+    #'max_autotune': True,
+    #'triton.cudagraphs': True,
+}
+
+def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
+
+    @spaces.GPU(duration=1500)
+    def compile_transformer():
+
+        with spaces.aoti_capture(pipeline.transformer) as call:
+            pipeline(*args, **kwargs)
+
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+        
+        exported = torch.export.export(
+            mod=pipeline.transformer,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+
+    spaces.aoti_apply(compile_transformer(), pipeline.transformer)

requirements.txt

@@ -0,0 +1,6 @@
+transformers
+accelerate
+safetensors
+bitsandbytes
+torchao
+kernels