Create app.py

app.py

@@ -0,0 +1,68 @@
+import os
+from pathlib import Path
+from typing import List
+import torch
+from PIL import Image
+import gradio as gr
+from ultraflux.pipeline_flux import FluxPipeline
+from ultraflux.transformer_flux import FluxTransformer2DModel
+from ultraflux.autoencoder_kl import AutoencoderKL
+
+torch.set_num_threads(os.cpu_count())
+torch.set_float32_matmul_precision("high")
+
+local_vae = AutoencoderKL.from_pretrained(
+    "Owen777/UltraFlux-v1",
+    subfolder="vae",
+    torch_dtype=torch.float32
+)
+
+transformer = FluxTransformer2DModel.from_pretrained(
+    "Owen777/UltraFlux-v1-1-Transformer",
+    torch_dtype=torch.float32
+)
+
+pipe = FluxPipeline.from_pretrained(
+    "Owen777/UltraFlux-v1",
+    vae=local_vae,
+    torch_dtype=torch.float32,
+    transformer=transformer
+)
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+pipe.scheduler.config.use_dynamic_shifting = False
+pipe.scheduler.config.time_shift = 4
+pipe = pipe.to("cpu")
+
+os.makedirs("results", exist_ok=True)
+
+def generate_ultraflux(prompt: str, seed: int = 0, steps: int = 50, size: int = 1024, guidance: float = 4.0):
+    out_path = Path("results") / f"ultra_flux.png"
+    with torch.inference_mode():
+        image = pipe(
+            prompt,
+            height=size,
+            width=size,
+            guidance_scale=guidance,
+            num_inference_steps=steps,
+            max_sequence_length=512,
+            generator=torch.Generator("cpu").manual_seed(seed)
+        ).images[0]
+    image.save(out_path)
+    return out_path
+
+demo = gr.Interface(
+    fn=generate_ultraflux,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+        gr.Number(label="Seed", value=0),
+        gr.Slider(10, 100, step=1, value=50, label="Inference Steps"),
+        gr.Slider(256, 2048, step=128, value=1024, label="Image Size"),
+        gr.Slider(1.0, 10.0, step=0.1, value=4.0, label="Guidance Scale")
+    ],
+    outputs=gr.Image(type="filepath"),
+    title="UltraFlux CPU Demo",
+    description="Generate high-quality images with UltraFlux on CPU."
+)
+
+demo.launch()