Update app.py

app.py

@@ -1,260 +1,131 @@
 import gradio as gr
-import numpy as np
-import random
-import os
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline, StableDiffusionPipeline
-from peft import PeftModel, LoraConfig
 import torch
-from typing import Optional
-
+from diffusers import StableDiffusionPipeline
+from peft import PeftModel, LoraConfig
+import os
 
 def get_lora_sd_pipeline(
     ckpt_dir='./lora_logos', 
     base_model_name_or_path=None, 
     dtype=torch.float16, 
     adapter_name="default"
-):
+    ):
     unet_sub_dir = os.path.join(ckpt_dir, "unet")
     text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
+    
     if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
         config = LoraConfig.from_pretrained(text_encoder_sub_dir)
         base_model_name_or_path = config.base_model_name_or_path
-
+    
     if base_model_name_or_path is None:
         raise ValueError("Please specify the base model name or path")
-
+    
     pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)
     pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
-
+    
     if os.path.exists(text_encoder_sub_dir):
-        pipe.text_encoder = PeftModel.from_pretrained(
-            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
-        )
-
+        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name)
+    
     if dtype in (torch.float16, torch.bfloat16):
         pipe.unet.half()
         pipe.text_encoder.half()
+    
     return pipe
 
-def split_prompt(prompt, tokenizer, max_length=77):
-    print(prompt)
-    print(type(prompt))
-    print(str(prompt))
-    tokens = tokenizer(prompt, truncation=False)["input_ids"]
-    chunks = [tokens[i:i + max_length] for i in range(0, len(tokens), max_length)]
-    return chunks
-
-def get_prompt_embeds(prompt_chunks, text_encoder):
-    prompt_embeds = []
-    for chunk in prompt_chunks:
-        chunk_tensor = torch.tensor([chunk]).to(text_encoder.device)
-        with torch.no_grad():
-            embeds = text_encoder(chunk_tensor)[0]
-        prompt_embeds.append(embeds)
-    return torch.cat(prompt_embeds, dim=1)
+def process_prompt(prompt, tokenizer, text_encoder, max_length=77):
+    tokens = tokenizer(prompt, truncation=False, return_tensors="pt")["input_ids"]
+    chunks = [tokens[:, i:i + max_length] for i in range(0, tokens.shape[1], max_length)]
+    
+    with torch.no_grad():
+        embeds = [text_encoder(chunk.to(text_encoder.device))[0] for chunk in chunks]
+    
+    return torch.cat(embeds, dim=1)
 
-def shape_alignment(prompt_embeds, negative_prompt_embeds):
+def align_embeddings(prompt_embeds, negative_prompt_embeds):
     max_length = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])
+    return torch.nn.functional.pad(prompt_embeds, (0, 0, 0, max_length - prompt_embeds.shape[1])), \
+           torch.nn.functional.pad(negative_prompt_embeds, (0, 0, 0, max_length - negative_prompt_embeds.shape[1]))
 
-    def pad_to_max_length(tensor, target_length):
-        padding = target_length - tensor.shape[1]
-        if padding > 0:
-            pad_tensor = torch.zeros(
-                tensor.shape[0], padding, tensor.shape[2], device=tensor.device
-            )
-            tensor = torch.cat([tensor, pad_tensor], dim=1)
-        return tensor
-
-    prompt_embeds = pad_to_max_length(prompt_embeds, max_length)
-    negative_prompt_embeds = pad_to_max_length(negative_prompt_embeds, max_length)
-
-    assert prompt_embeds.shape == negative_prompt_embeds.shape, "Shapes do not match!"
-    return prompt_embeds, negative_prompt_embeds
-
-def prompts_embeddings(prompt, negative_promt, tokenizer, text_encoder):
-    prompt_chunks = split_prompt(prompt, tokenizer)
-    negative_prompt_chunks = split_prompt(negative_prompt, tokenizer)
-
-    prompt_embeds = get_prompt_embeds(prompt_chunks, text_encoder)
-    negative_prompt_embeds = get_prompt_embeds(negative_prompt_chunks, text_encoder)
-
-    prompt_embeds, negative_prompt_embeds = shape_alignment(prompt_embeds, negative_prompt_embeds)
-
-    return prompt_embeds, negative_prompt_embeds
-
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_id_default = "CompVis/stable-diffusion-v1-4"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-
-pipe_default = get_lora_sd_pipeline(
-    ckpt_dir='./lora_logos', 
-    base_model_name_or_path=model_id_default, 
-    dtype=torch_dtype,
-    )
-# pipe_default = DiffusionPipeline.from_pretrained(model_id_default, torch_dtype=torch_dtype)
-pipe_default = pipe_default.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+pipe_default = get_lora_sd_pipeline(ckpt_dir='./lora_logos', base_model_name_or_path=model_id_default, dtype=torch_dtype).to(device)
 
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
-    prompt: str,
-    negative_prompt: str,
-    width: int,
-    height: int,
-    num_inference_steps: Optional[int] = 20,
-    model_id: Optional[str] = 'CompVis/stable-diffusion-v1-4',
-    seed: Optional[int] = 42,
-    guidance_scale: Optional[float] = 7.0,
-    lora_scale: Optional[float] = 0.5,
-    progress=gr.Progress(track_tqdm=True),
-):
-    generator = torch.Generator().manual_seed(seed)
+    prompt, 
+    negative_prompt, 
+    width=512, 
+    height=512, 
+    num_inference_steps=20, 
+    model_id='CompVis/stable-diffusion-v1-4', 
+    seed=42, 
+    guidance_scale=7.0, 
+    lora_scale=0.5
+    ):
+    generator = torch.Generator(device).manual_seed(seed)
+
+    print(prompt)
+    print(type(prompt))
 
+    print(negative_prompt)
+    print(type(negative_prompt))
+    
+    if model_id != model_id_default:
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype).to(device)
+        prompt_embeds = process_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
+        negative_prompt_embeds = process_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+    else:
+        pipe = pipe_default
+        prompt_embeds = process_prompt(prompt, pipe.tokenizer, pipe.text_encoder)
+        negative_prompt_embeds = process_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+        pipe.fuse_lora(lora_scale=lora_scale)
+    
     params = {
-        # 'prompt': prompt,
-        # 'negative_prompt': negative_prompt,
+        'prompt_embeds': prompt_embeds,
+        'negative_prompt_embeds': negative_prompt_embeds,
         'guidance_scale': guidance_scale,
         'num_inference_steps': num_inference_steps,
         'width': width,
         'height': height,
         'generator': generator,
     }
-
-    if model_id != model_id_default:
-        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
-        pipe = pipe.to(device)
-        image = pipe(**params).images[0]
-    else:
-        prompt_embeds, negative_prompt_embeds = prompts_embeddings(
-            prompt, 
-            negative_prompt, 
-            pipe_default.tokenizer, 
-            pipe_default.text_encoder
-            )
-        params['prompt_embeds'] = prompt_embeds
-        params['negative_prompt_embeds']=negative_prompt_embeds
-        pipe_default.fuse_lora(lora_scale=lora_scale)
-        image = pipe_default(**params).images[0]
-
-    return image
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
+    
+    return pipe(**params).images[0]
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        gr.Markdown("# DEMO Text-to-Image")
+        model_id = gr.Textbox(label="Model ID", value=model_id_default)
+        prompt = gr.Textbox(label="Prompt")
+        negative_prompt = gr.Textbox(label="Negative prompt")
+        seed = gr.Number(label="Seed", value=42)
+        guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, value=7.0)
+        lora_scale = gr.Slider(label="LoRA scale", minimum=0.0, maximum=1.0, value=0.5)
+        num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, value=20)
         
-        gr.Markdown(" # DEMO Text-to-Image")
-
-        with gr.Row():
-            model_id = gr.Textbox(
-                label="Model ID",
-                max_lines=1,
-                placeholder="Enter model id like 'CompVis/stable-diffusion-v1-4'",
-                value="CompVis/stable-diffusion-v1-4"
-            )
-
-        prompt = gr.Textbox(
-            label="Prompt",
-            max_lines=1,
-            placeholder="Enter your prompt",
-        )
-
-        negative_prompt = gr.Textbox(
-            label="Negative prompt",
-            max_lines=1,
-            placeholder="Enter a negative prompt",
-        )
-
-        with gr.Row():
-            seed = gr.Number(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=42,
-            )
-
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance scale",
-                minimum=0.0,
-                maximum=10.0,
-                step=0.1,
-                value=7.0,
-            )
-
-        with gr.Row():
-            lora_scale = gr.Slider(
-                label="LoRA scale",
-                minimum=0.0,
-                maximum=1.0,
-                step=0.1,
-                value=0.5,
-            )
-
-        with gr.Row():
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps",
-                minimum=1,
-                maximum=50,
-                step=1,
-                value=20,
-            )
-
         with gr.Accordion("Optional Settings", open=False):
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-            
-            with gr.Row():
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-
-        run_button = gr.Button("Run", scale=1, variant="primary")
-        result = gr.Image(label="Result", show_label=False)
+            width = gr.Slider(label="Width", minimum=256, maximum=1024, value=512, step=32)
+            height = gr.Slider(label="Height", minimum=256, maximum=1024, value=512, step=32)
+        
+        run_button = gr.Button("Run")
+        result = gr.Image(label="Result")
     
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+    run_button.click(
+        fn=infer, 
         inputs=[
-            prompt,
-            negative_prompt,
-            width,
-            height,
-            num_inference_steps,
-            model_id,
-            seed,
-            guidance_scale,
-            lora_scale,
-        ],
-        outputs=[result],
-    )
+            prompt, 
+            negative_prompt, 
+            width, 
+            height, 
+            num_inference_steps, 
+            model_id, seed, 
+            guidance_scale, 
+            lora_scale
+        ], 
+        outputs=result)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()