Update app.py

app.py

@@ -1,23 +1,22 @@
 import os
 import gradio as gr
 import torch
-from diffusers import StableAudioPipeline
-import spaces
-from translatepy import Translator
 import numpy as np
 import random
-import soundfile as sf
+from diffusers import FluxPipeline
+import spaces
+import transformers
+from translatepy import Translator
 
-translator = Translator()
 
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+translator = Translator()
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Constants
-model = "stabilityai/stable-audio-open-1.0"
+model = "black-forest-labs/FLUX.1-dev"
 MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
-.gradio-container {
-  max-width: 690px !important;
-}
 footer {
     visibility: hidden;
 }
@@ -29,73 +28,135 @@ JS = """function () {
     window.location.replace(gradioURL + '?__theme=dark');
   }
 }"""
-DESCRIPTION = """
-<center>
-Stable Audio Open 1.0 generates variable-length (up to 47s) stereo audio at 44.1kHz from text prompts. \
-It comprises three components: an autoencoder that compresses waveforms into a manageable sequence length, \
-a T5-based text embedding for text conditioning, and a transformer-based diffusion (DiT) model that operates in the latent space of the autoencoder.
-</center>
-"""
-
-pipe = StableAudioPipeline.from_pretrained(
-    model,
-    torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
 
+if torch.cuda.is_available():
+    pipe = FluxPipeline.from_pretrained(
+        model, 
+        torch_dtype=torch.bfloat16)
 
 # Function 
-@spaces.GPU(duration=120)
-def main(
+@spaces.GPU()
+def generate_image(
     prompt,
-    negative="low quality",
-    second: float = 10.0,
-    seed: int = -1):
-    
+    width=1024,
+    height=1024,
+    scales=5,
+    steps=30,
+    seed: int =-1,
+    nums=1,
+    progress=gr.Progress(track_tqdm=True)):
+
     if seed == -1:
-         seed = random.randint(0, MAX_SEED)
+        seed = random.randint(0, MAX_SEED)
     seed = int(seed)
-    generator = torch.Generator().manual_seed(seed)
+    print(f'prompt:{prompt}')
     
-    prompt = str(translator.translate(prompt, 'English'))
+    text = str(translator.translate(prompt['text'], 'English'))
 
-    print(f'prompt:{prompt}')
+    
+    generator = torch.Generator().manual_seed(seed)
+    
 
-    audio = pipe(
-        prompt, 
-        negative_prompt=negative, 
-        audio_end_in_s=second,
-        num_inference_steps=200,
-        num_waveforms_per_prompt=3,
+    image = pipe(
+        prompt,
+        height=height,
+        width=width,
+        guidance_scale=scale,
+        output_type="pil",
+        num_inference_steps=steps,
+        max_sequence_length=512,
+        num_images_per_prompt=nums,
         generator=generator,
-    ).audios
+    ).images
 
-    os.makedirs("outputs", exist_ok=True)
-    base_count = len(glob(os.path.join("outputs", "*.mp4")))
-    audio_path = os.path.join("outputs", f"{base_count:06d}.wav")
-    
-    sf.write(audio_path, audio[0].T.float().cpu().numpy(), pipe.vae.samping_rate)
-    
-    return audio_path, seed
+    print(image)
+    print(seed)
+    return image, seed
 
-# Gradio Interface
 
-with gr.Blocks(theme='soft', css=CSS, js=JS, title="Stable Audio Open") as iface:
-    with gr.Accordion(""):
-        gr.Markdown(DESCRIPTION)
-    output = gr.Audio(label="Podcast", type="filepath", interactive=False, autoplay=True, elem_classes="audio")  # Create an output textbox
-    prompt = gr.Textbox(label="Prompt", placeholder="1000 BPM percussive sound of water drops")
-    negative = gr.Textbox(label="Negative prompt", placeholder="Low quality")
-    with gr.Row():
-        second =  gr.Slider(5.0, 60.0, value=10.0, label="Second", step=0.1),
-        seed = gr.Slider(-1, MAX_SEED, value=-1, label="Seed", step=1),
-    with gr.Row():
-        submit_btn = gr.Button("🚀 Send")  # Create a submit button
-        clear_btn = gr.ClearButton([prompt, seed, output], value="🗑️ Clear") # Create a clear button
+examples = [
+        "a female character with long, flowing hair that appears to be made of ethereal, swirling patterns resembling the Northern Lights or Aurora Borealis. The background is dominated by deep blues and purples, creating a mysterious and dramatic atmosphere. The character's face is serene, with pale skin and striking features. She wears a dark-colored outfit with subtle patterns. The overall style of the artwork is reminiscent of fantasy or supernatural genres",
+        "Digital art, portrait of an anthropomorphic roaring Tiger warrior with full armor, close up in the middle of a battle, behind him there is a banner with the text \"Open Source\".", 
+        "photo of a dog and a cat both standing on a red box, with a blue ball in the middle with a parrot standing on top of the ball. The box has the text \"FLUX\"", 
+		"selfie photo of a wizard with long beard and purple robes, he is apparently in the middle of Tokyo. Probably taken from a phone.", 
+		"A vibrant street wall covered in colorful graffiti, the centerpiece spells \"FLUX\", in a storm of colors",
+		"photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.", 
+		"anime art of a steampunk inventor in their workshop, surrounded by gears, gadgets, and steam. He is holding a blue potion and a red potion, one in each hand", 
+		"photo of picturesque scene of a road surrounded by lush green trees and shrubs. The road is wide and smooth, leading into the distance. On the right side of the road, there's a blue sports car parked with the license plate spelling \"FLUX\". The sky above is partly cloudy, suggesting a pleasant day. The trees have a mix of green and brown foliage. There are no people visible in the image. The overall composition is balanced, with the car serving as a focal point.",
+		"photo of young man in a black suit, white shirt, and black tie. He has a neatly styled haircut and is looking directly at the camera with a neutral expression. The background consists of a textured wall with horizontal lines. The photograph is in black and white, emphasizing contrasts and shadows. The man appears to be in his late twenties or early thirties, with fair skin and short, dark hair.",
+		"photo of a woman on the beach, shot from above. She is facing the sea, while wearing a white dress. She has long blonde hair"
+]
+
 
-    # Set up the event listeners
-    submit_btn.click(main, inputs=[prompt, negative, second, seed], outputs=[output, seed])
 
+# Gradio Interface
 
-#gr.close_all()
+with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
+    gr.HTML("<h1><center>Flux Dev</center></h1>")
+    gr.HTML("<p><center><a href='https://huggingface.co/black-forest-labs/FLUX.1-dev'>FLUX</a> text-to-image generation<br><b>Feature</b>: Support 512 token and multi-languages</center></p>")
+    with gr.Row():
+        with gr.Column(scale=4):
+            img = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
+            prompt = gr.Textbox(label='Enter Your Prompt (Multi-Languages)', placeholder="Enter prompt...")
+        with gr.Accordion("Advanced Options", open=True):
+            with gr.Column(scale=1):
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=1280,
+                    step=8,
+                    value=1024,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=1280,
+                    step=8,
+                    value=1024,
+                )
+                scales = gr.Slider(
+                    label="Guidance",
+                    minimum=3.5,
+                    maximum=7,
+                    step=0.1,
+                    value=3.5,
+                )
+                steps = gr.Slider(
+                    label="Steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=30,
+                )
+                seed = gr.Slider(
+                    label="Seed (-1 Random)",
+                    minimum=-1,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=-1,
+                    scale=2,
+                )
+                nums = gr.Slider(
+                    label="Image Numbers",
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    value=1,
+                    scale=1,
+                )  
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=[img, seed],
+        fn=generate_image,
+        cache_examples="lazy",
+        examples_per_page=4,
+    )
+
+    prompt.submit(fn=generate_image,
+                 inputs=[prompt, width, height, scales, steps, seed, nums],
+                 outputs=[img, seed],
+                 )
 
-iface.queue().launch(show_api=False)  # Launch the Gradio interface
+    
+demo.queue().launch()