up

app-fast.py

@@ -36,13 +36,10 @@ RESOLUTION_OPTIONS: list[str] = [
     "832 x 1248",
 ]
 
-device = torch.device("cuda")
-
 quant_config = TransformersBitsAndBytesConfig(
     load_in_4bit=True,
 )
 
-
 tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
@@ -71,8 +68,8 @@ scheduler = MODEL_CONFIGS["scheduler"](
 
 pipe = HiDreamImagePipeline.from_pretrained(
     MODEL_PATH,
-    scheduler=scheduler,
     transformer=transformer,
+    scheduler=scheduler,
     tokenizer_4=tokenizer,
     text_encoder_4=text_encoder,
     device_map="balanced",
@@ -90,8 +87,6 @@ def generate_image(
     if seed == -1:
         seed = torch.randint(0, 1_000_000, (1,)).item()
 
-    # msg = "ℹ️ This spaces currently crash because of the memory usage. Please help me fix 😅"
-    # raise gr.Error(msg, duration=10)
     height, width = tuple(map(int, resolution.replace(" ", "").split("x")))
     generator = torch.Generator("cuda").manual_seed(seed)
 
@@ -128,11 +123,6 @@ with gr.Blocks(title="HiDream Image Generator Fast") as demo:
 
             seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
             generate_btn = gr.Button("Generate Image", variant="primary")
-            # generate_btn = gr.Button(
-            #     "This space currently crash because of the memory usage. Please help me fix 😅",
-            #     variant="primary",
-            #     interactive=False,
-            # )
             seed_used = gr.Number(label="Seed Used", interactive=False)
 
         with gr.Column():