Update app.py

app.py

@@ -1,31 +1,24 @@
-import torch
-from diffusers import FluxPipeline
 import gradio as gr
+from diffusers import DiffusionPipeline
+import torch
 
-# Initialize the model
-pipe = FluxPipeline.from_pretrained("Shakker-Labs/AWPortrait-FL", torch_dtype=torch.bfloat16)
-pipe.to("cuda")
+# Load the Kwai-Kolors diffusion model
+pipeline = DiffusionPipeline.from_pretrained("Kwai-Kolors/Kolors").to("cuda" if torch.cuda.is_available() else "cpu")
 
+# Function to generate image from prompt
 def generate_image(prompt):
-    # Generate the image
-    image = pipe(prompt, 
-                 num_inference_steps=24, 
-                 guidance_scale=3.5,
-                 width=768, height=1024,
-                ).images[0]
-    # Save the image
-    image_path = "generated_image.png"
-    image.save(image_path)
-    return image_path
+    # Use the pipeline to generate an image from the text prompt
+    image = pipeline(prompt).images[0]
+    return image
 
-# Define the Gradio interface
-interface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
-    outputs=gr.Image(type="file", label="Generated Image"),
-    title="Image Generator",
-    description="Generate images based on the given prompt using the FluxPipeline model."
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,                # Function that generates the image
+    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt"),  # Textbox input for the prompt
+    outputs="image",                  # Output is an image
+    title="Kwai-Kolors Image Generator",
+    description="Generate images from text prompts using the Kwai-Kolors diffusion model."
 )
 
-# Launch the Gradio app
-interface.launch()
+# Launch the app
+iface.launch()