Update app.py

app.py

@@ -1,85 +1,27 @@
 import os
-import torch
 import random
-from huggingface_hub import snapshot_download
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline
-from kolors.models.modeling_chatglm import ChatGLMModel
-from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from diffusers import UNet2DConditionModel, AutoencoderKL
-from diffusers import EulerDiscreteScheduler
 import gradio as gr
+import torch
+from diffusers import DiffusionPipeline
 
-# Download the model files
-ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
-
-# Function to load models
-def load_models():
-    text_encoder = ChatGLMModel.from_pretrained(
-        os.path.join(ckpt_dir, 'text_encoder'),
-        torch_dtype=torch.float16).half()
-    tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(ckpt_dir, 'text_encoder'))
-    vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), revision=None).half()
-    scheduler = EulerDiscreteScheduler.from_pretrained(os.path.join(ckpt_dir, "scheduler"))
-    unet = UNet2DConditionModel.from_pretrained(os.path.join(ckpt_dir, "unet"), revision=None).half()
-
-    return StableDiffusionXLPipeline(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        scheduler=scheduler,
-        force_zeros_for_empty_prompt=False
-    )
-
-# Create a global variable to hold the pipeline
-pipe = load_models()
+pipe = DiffusionPipeline.from_pretrained("fluently/Fluently-XL-Final")
+pipe.load_lora_weights("OEvortex/HelpingAI-PixelCraft")
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def generate_image(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, num_images_per_prompt, use_random_seed, seed, progress=gr.Progress(track_tqdm=True)):
-    if use_random_seed:
-        seed = random.randint(0, 2**32 - 1)
-    else:
-        seed = int(seed)  # Ensure seed is an integer
-    
-    # Move the model to the CPU for inference and clear unnecessary variables
-    with torch.no_grad():
-        generator = torch.Generator().manual_seed(seed)
-        result = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images_per_prompt,
-            generator=generator
-        )
-        image = result.images
-    
-    return image, seed
+# Function to generate image from prompt
+def generate_image(prompt):
+    # Use the pipeline to generate an image from the text prompt
+    image = pipe(prompt).images[0]
+    return image
 
-# Gradio interface
+# Create Gradio interface
 iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Textbox(label="Negative Prompt")
-    ],
-    additional_inputs=[
-        gr.Slider(512, 2048, 1024, step=64, label="Height"),
-        gr.Slider(512, 2048, 1024, step=64, label="Width"),
-        gr.Slider(20, 50, 20, step=1, label="Number of Inference Steps"),
-        gr.Slider(1, 20, 5, step=0.5, label="Guidance Scale"),
-        gr.Slider(1, 4, 1, step=1, label="Number of images per prompt"),
-        gr.Checkbox(label="Use Random Seed", value=True),
-        gr.Number(label="Seed", value=0, precision=0)
-    ],
-    additional_inputs_accordion=gr.Accordion(label="Advanced settings", open=False),
-    outputs=[
-        gr.Gallery(label="Result", elem_id="gallery", show_label=False),
-        gr.Number(label="Seed Used")
-    ],
-    title="Kolors",
-    theme='bethecloud/storj_theme',
+    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."
 )
 
-iface.launch()
+# Launch the app
+iface.launch()