Add App

app.py

@@ -0,0 +1,123 @@
+from datasets import load_dataset
+import gradio as gr
+from gradio_client import Client
+import json
+import torch
+from diffusers import FluxPipeline, AutoencoderKL
+from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+import spaces
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+# pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
+# pipe.to(torch.float16)
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+
+llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
+# t2i_client = Client("black-forest-labs/FLUX.1-dev")
+# t2i_client = Client("black-forest-labs/FLUX.1-schnell")
+
+ds = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")
+
+prompt_template = """Generate a character with this persona description: {persona_description}
+In a world with this description: {world_description}
+
+Write the character in json format with the following fields:
+- name: The name of the character
+- background: The background of the character
+- appearance: The appearance of the character
+- personality: The personality of the character
+- skills_and_abilities: The skills and abilities of the character
+- goals: The goals of the character
+- conflicts: The conflicts of the character
+- backstory: The backstory of the character
+- current_situation: The current situation of the character
+- spoken_lines: The spoken lines of the character (list of strings)
+
+Don't write anything else except the character description in json format and don't include '```'.
+"""
+
+world_description_prompt = "Generate a unique and random world description (Don't Write anything else except the world description)."
+
+def get_random_world_description():
+    result = llm_client.predict(
+        query=world_description_prompt,
+        history=[],
+        system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+        api_name="/model_chat",
+    )
+    return result[1][0][-1]
+
+def get_random_persona_description():
+    return ds.shuffle().select([100])[0]["persona"]
+
+@spaces.GPU()
+def generate_character(world_description, persona_description, progress=gr.Progress(track_tqdm=True)):
+    result = llm_client.predict(
+        query=prompt_template.format(
+            persona_description=persona_description, world_description=world_description
+        ),
+        history=[],
+        system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+        api_name="/model_chat",
+    )
+    output = json.loads(result[1][0][-1])
+    print("Character generated")
+    print(json.dumps(output, indent=4))
+
+    for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        prompt=output["appearance"],
+        guidance_scale=3.5,
+        num_inference_steps=28,
+        width=1024,
+        height=1024,
+        generator=torch.Generator("cpu").manual_seed(0),
+        output_type="pil",
+        good_vae=good_vae,
+    ):
+        yield image, output
+    print("Character and image generated")
+
+
+with gr.Blocks(title="Character Generator") as app:
+    with gr.Column():
+        gr.HTML("<center><h1>Character Generator</h1></center>")
+        gr.HTML(
+            "<center><h3>Generate a character with a persona description and a world description.</h3></center>"
+        )
+        with gr.Column():
+            with gr.Row():
+                world_description = gr.Textbox(lines=10, label="World Description", scale=4)
+                persona_description = gr.Textbox(lines=10, label="Persona Description", value=get_random_persona_description(), scale=1)
+            with gr.Row():
+                random_world_button = gr.Button(value="Get Random World Description", variant="secondary", scale=1)
+                submit_button = gr.Button(value="Generate Interesting Character!", variant="primary", scale=5)
+                random_persona_button = gr.Button(value="Get Random Persona Description", variant="secondary", scale=1)
+        with gr.Row():
+            character_image = gr.Image(label="Character Image")
+            character = gr.JSON(label="Character Description")
+
+        examples = gr.Examples(
+            [
+                "In a world where magic is real and dragons roam the skies, a group of adventurers set out to find the legendary sword of the dragon king.",
+                "Welcome to Aethoria, a vast and mysterious realm where the laws of physics bend to the will of ancient magic. This world is comprised of countless floating islands suspended in an endless sky, each one a unique ecosystem teeming with life and secrets. The islands of Aethoria range from lush, verdant jungles to barren, crystalline deserts. Some are no larger than a city block, while others span hundreds of miles. Connecting these disparate landmasses are shimmering bridges of pure energy, and those brave enough to venture off the beaten path can find hidden portals that instantly transport them across great distances. Aethoria's inhabitants are as diverse as its landscapes. Humans coexist with ethereal beings of light, rock-skinned giants, and shapeshifting creatures that defy classification. Ancient ruins dot the islands, hinting at long-lost civilizations and forgotten technologies that blur the line between science and sorcery. The world is powered by Aether, a mystical substance that flows through everything. Those who can harness its power become formidable mages, capable of manipulating reality itself. However, Aether is a finite resource, and its scarcity has led to conflicts between the various factions vying for control. In the skies between the islands, magnificent airships sail on currents of magic, facilitating trade and exploration. Pirates and sky raiders lurk in the cloudy depths, always on the lookout for unsuspecting prey. Deep beneath the floating lands lies the Undervoid, a dark and treacherous realm filled with nightmarish creatures and untold riches. Only the bravest adventurers dare to plumb its depths, and fewer still return to tell the tale. As an ever-present threat, the Chaos Storms rage at the edges of the known world, threatening to consume everything in their path. It falls to the heroes of Aethoria to uncover the secrets of their world and find a way to push back the encroaching darkness before it's too late. In Aethoria, every island holds a story, every creature has a secret, and every adventure could change the fate of this wondrous, imperiled world.",
+                "In a world from my imagination, there is a city called 'Orakis'. floating in the sky on pillars of pure light. The walls of the city are made of crystal glass, constantly reflecting the colors of dawn and dusk, giving it an eternal celestial glow. The buildings breathe and change their shapes according to the seasons—they grow in spring, strengthen in summer, and begin to fade in autumn until they become mist in winter.",
+            ],
+            world_description,
+        )
+
+    submit_button.click(
+        generate_character, [world_description, persona_description], outputs=[character_image, character]
+    )
+    random_world_button.click(
+        get_random_world_description, outputs=[world_description]
+    )
+    random_persona_button.click(
+        get_random_persona_description, outputs=[persona_description]
+    )
+
+app.queue().launch(share=True)

live_preview_helpers.py

@@ -0,0 +1,166 @@
+import torch
+import numpy as np
+from diffusers import FluxPipeline, AutoencoderTiny, FlowMatchEulerDiscreteScheduler
+from typing import Any, Dict, List, Optional, Union
+
+# Helper functions
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+# FLUX pipeline function
+@torch.inference_mode()
+def flux_pipe_call_that_returns_an_iterable_of_images(
+    self,
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 28,
+    timesteps: List[int] = None,
+    guidance_scale: float = 3.5,
+    num_images_per_prompt: Optional[int] = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.FloatTensor] = None,
+    prompt_embeds: Optional[torch.FloatTensor] = None,
+    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    max_sequence_length: int = 512,
+    good_vae: Optional[Any] = None,
+):
+    height = height or self.default_sample_size * self.vae_scale_factor
+    width = width or self.default_sample_size * self.vae_scale_factor
+
+    # 1. Check inputs
+    self.check_inputs(
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        max_sequence_length=max_sequence_length,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._joint_attention_kwargs = joint_attention_kwargs
+    self._interrupt = False
+
+    # 2. Define call parameters
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    device = self._execution_device
+
+    # 3. Encode prompt
+    lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
+    prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
+        prompt=prompt,
+        prompt_2=prompt_2,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        max_sequence_length=max_sequence_length,
+        lora_scale=lora_scale,
+    )
+    # 4. Prepare latent variables
+    num_channels_latents = self.transformer.config.in_channels // 4
+    latents, latent_image_ids = self.prepare_latents(
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    image_seq_len = latents.shape[1]
+    mu = calculate_shift(
+        image_seq_len,
+        self.scheduler.config.base_image_seq_len,
+        self.scheduler.config.max_image_seq_len,
+        self.scheduler.config.base_shift,
+        self.scheduler.config.max_shift,
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        self.scheduler,
+        num_inference_steps,
+        device,
+        timesteps,
+        sigmas,
+        mu=mu,
+    )
+    self._num_timesteps = len(timesteps)
+
+    # Handle guidance
+    guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
+
+    # 6. Denoising loop
+    for i, t in enumerate(timesteps):
+        if self.interrupt:
+            continue
+
+        timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+        noise_pred = self.transformer(
+            hidden_states=latents,
+            timestep=timestep / 1000,
+            guidance=guidance,
+            pooled_projections=pooled_prompt_embeds,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=text_ids,
+            img_ids=latent_image_ids,
+            joint_attention_kwargs=self.joint_attention_kwargs,
+            return_dict=False,
+        )[0]
+        # Yield intermediate result
+        latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents_for_image, return_dict=False)[0]
+        yield self.image_processor.postprocess(image, output_type=output_type)[0]
+        
+        latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        torch.cuda.empty_cache()
+
+    # Final image using good_vae
+    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+    latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
+    image = good_vae.decode(latents, return_dict=False)[0]
+    self.maybe_free_model_hooks()
+    torch.cuda.empty_cache()
+    yield self.image_processor.postprocess(image, output_type=output_type)[0]

requirements.txt

@@ -0,0 +1,6 @@
+spaces
+torch
+diffusers
+gradio_client
+datasets
+numpy